[glsl-in] Implement swizzle for r-values

Related to gfx-rs#210
pjoe · Sep 24, 2020 · b78698d · b78698d
1 parent 7ce9977
commit b78698d
Show file tree

Hide file tree

Showing 2 changed files with 90 additions and 14 deletions.
diff --git a/src/front/glsl/parser.rs b/src/front/glsl/parser.rs
@@ -218,20 +218,9 @@ pomelo! {
     }
     postfix_expression ::= function_call;
     postfix_expression ::= postfix_expression(e) Dot Identifier(i) /* FieldSelection in spec */ {
-        let type_inner = extra.resolve_type(e.expression)?;
-        if let TypeInner::Struct{members} = type_inner {
-            let index = members.iter().position(|m| m.name == Some(i.1.clone()))
-                .ok_or(ErrorKind::UnknownField(i.0, i.1))?;
-            ExpressionRule::from_expression(
-                extra.context.expressions.append(Expression::AccessIndex {
-                    base: e.expression,
-                    index: index as u32,
-                })
-            )
-        } else {
-            //TODO: swizzle
-            return Err(ErrorKind::SemanticError("Can't lookup field on this type"))
-        }
+        //TODO: how will this work as l-value?
+        let exp = extra.field_selection(e.expression, &*i.1, i.0)?;
+        ExpressionRule::from_expression(exp.ok_or(ErrorKind::SemanticError("Can't lookup field on this type"))?)
     }
     postfix_expression ::= postfix_expression(pe) IncOp {
         //TODO

diff --git a/src/front/glsl/variables.rs b/src/front/glsl/variables.rs
@@ -5,6 +5,7 @@ use crate::{
 
 use super::ast::*;
 use super::error::ErrorKind;
+use super::token::TokenMetadata;
 
 impl Program {
     pub fn lookup_variable(&mut self, name: &str) -> Result<Option<Handle<Expression>>, ErrorKind> {
@@ -96,4 +97,90 @@ impl Program {
             Ok(None)
         }
     }
+
+    pub fn field_selection(
+        &mut self,
+        expression: Handle<Expression>,
+        name: &str,
+        meta: TokenMetadata,
+    ) -> Result<Option<Handle<Expression>>, ErrorKind> {
+        let type_inner = self.resolve_type(expression)?;
+        match type_inner {
+            TypeInner::Struct { members } => {
+                let index = members
+                    .iter()
+                    .position(|m| m.name == Some(name.into()))
+                    .ok_or_else(|| ErrorKind::UnknownField(meta, name.into()))?;
+                Ok(Some(self.context.expressions.append(
+                    Expression::AccessIndex {
+                        base: expression,
+                        index: index as u32,
+                    },
+                )))
+            }
+            // swizzles (xyzw, rgba, stpq)
+            &TypeInner::Vector { size, kind, width } => {
+                let check_swizzle_components = |comps: &str| {
+                    name.chars()
+                        .map(|c| {
+                            comps
+                                .find(c)
+                                .and_then(|i| if i < size as usize { Some(i) } else { None })
+                        })
+                        .fold(Some(Vec::<usize>::new()), |acc, cur| {
+                            cur.and_then(|i| {
+                                acc.map(|mut v| {
+                                    v.push(i);
+                                    v
+                                })
+                            })
+                        })
+                };
+
+                let indices = check_swizzle_components("xyzw")
+                    .or_else(|| check_swizzle_components("rgba"))
+                    .or_else(|| check_swizzle_components("stpq"));
+
+                if let Some(v) = indices {
+                    let mut components: Vec<Handle<Expression>> = v
+                        .iter()
+                        .map(|idx| {
+                            self.context.expressions.append(Expression::AccessIndex {
+                                base: expression,
+                                index: *idx as u32,
+                            })
+                        })
+                        .collect();
+                    if components.len() == 1 {
+                        // only single element swizzle, like pos.y, just return that component
+                        Ok(components.pop())
+                    } else {
+                        Ok(Some(self.context.expressions.append(Expression::Compose {
+                            ty: self.module.types.fetch_or_append(Type {
+                                name: None,
+                                inner: TypeInner::Vector {
+                                    kind,
+                                    width,
+                                    size: match components.len() {
+                                        2 => VectorSize::Bi,
+                                        3 => VectorSize::Tri,
+                                        4 => VectorSize::Quad,
+                                        _ => {
+                                            return Err(ErrorKind::SemanticError(
+                                                "Bad swizzle size",
+                                            ));
+                                        }
+                                    },
+                                },
+                            }),
+                            components,
+                        })))
+                    }
+                } else {
+                    Ok(None)
+                }
+            }
+            _ => Ok(None),
+        }
+    }
 }