From df9834258f819ca01154252ceba5eef63b629734 Mon Sep 17 00:00:00 2001
From: Pelle Johnsen <pelle.johnsen@gmail.com>
Date: Fri, 25 Sep 2020 16:32:50 +0200
Subject: [PATCH] [glsl-in] Implement swizzle for r-values (#213)

* [glsl-in] Implement swizzle for r-values

Related to #210

* [glsl-in] Just return Result from field_selection

Removed unneccessary Otion in return type

* [glsl-in] Always match on type in field_selection

* [glsl.in] Borrow by value in field_selection
---
 src/front/glsl/parser.rs    | 17 ++------
 src/front/glsl/variables.rs | 84 +++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 14 deletions(-)

diff --git a/src/front/glsl/parser.rs b/src/front/glsl/parser.rs
index 31aab6d7be..d5da8f3850 100644
--- 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 expression = extra.field_selection(e.expression, &*i.1, i.0)?;
+        ExpressionRule { expression, statements: e.statements }
     }
     postfix_expression ::= postfix_expression(pe) IncOp {
         //TODO
diff --git a/src/front/glsl/variables.rs b/src/front/glsl/variables.rs
index 834f787651..c4d37b7bb3 100644
--- 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,87 @@ impl Program {
             Ok(None)
         }
     }
+
+    pub fn field_selection(
+        &mut self,
+        expression: Handle<Expression>,
+        name: &str,
+        meta: TokenMetadata,
+    ) -> Result<Handle<Expression>, ErrorKind> {
+        match *self.resolve_type(expression)? {
+            TypeInner::Struct { ref members } => {
+                let index = members
+                    .iter()
+                    .position(|m| m.name == Some(name.into()))
+                    .ok_or_else(|| ErrorKind::UnknownField(meta, name.into()))?;
+                Ok(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 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[0])
+                    } else {
+                        Ok(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 {
+                    Err(ErrorKind::SemanticError("Invalid swizzle for vector"))
+                }
+            }
+            _ => Err(ErrorKind::SemanticError("Can't lookup field on this type")),
+        }
+    }
 }