Add support for swizzling triple types.

Author: marsupial

src/liboslcomp/ast.cpp

@@ -719,6 +719,65 @@ ASTstructselect::print (std::ostream &out, int indentlevel) const
 
 
 
+ASTswizzle::ASTswizzle (OSLCompilerImpl *comp, ASTNode *expr, ustring field) :
+ASTfieldselect(swizzle_node, comp, expr, field)
+{
+    if (field.size() == 1) m_typespec = TypeDesc::TypeFloat;
+    else m_typespec = TypeDesc::TypeVector;
+}
+
+
+
+size_t ASTswizzle::indices (ustring components, int *indexes, size_t N, bool consts)
+{
+    size_t i = 0;
+    do {
+        switch (components[i]) {
+            case 'r':
+            case 'x':  indexes[i] = 0; break;
+            case 'g':
+            case 'y':  indexes[i] = 1; break;
+            case 'b':
+            case 'z':  indexes[i] = 2; break;
+
+            case '0':
+                if (!consts)
+                    return 0;
+                indexes[i] = const_offset;
+                break;
+            case '1':
+                if (!consts)
+                    return 0;
+                indexes[i] = const_offset + 1;
+                break;
+
+            case '\0': return i;
+
+            default:  return 0;
+        }
+    } while (++i < N);
+
+    return i;
+}
+
+
+
+const char *
+ASTswizzle::childname (size_t i) const
+{
+    return type_c_str(m_typespec);
+}
+
+
+void
+ASTswizzle::print (std::ostream &out, int indentlevel) const
+{
+    ASTNode::print (out, indentlevel);
+    indent (out, indentlevel+1);
+    out << "components " << field() << "\n";
+}
+
+
 ASTNode* ASTfieldselect::create (OSLCompilerImpl *comp, ASTNode *expr,
                                  ustring field)
 {
@@ -728,25 +787,55 @@ ASTNode* ASTfieldselect::create (OSLCompilerImpl *comp, ASTNode *expr,
     const TypeSpec &type = expr->nodetype() != structselect_node ? expr->typespec() :
                            static_cast<ASTstructselect*>(expr)->fieldsym()->typespec();
 
-    if (type.aggregate() == TypeDesc::VEC3) {
-        int component = -1;
-        switch (field[0]) {
-            case 'r': component = 0; break;
-            case 'g': component = 1; break;
-            case 'b': component = 2; break;
+    if (type.aggregate() == TypeDesc::VEC3 && field.size() <= 3) {
+        // Early out swizle to native component ordering.
+        if (field == "rgb" || field == "xyz")
+            return expr;
 
-            case 'x': component = 0; break;
-            case 'y': component = 1; break;
-            case 'z': component = 2; break;
+        ASTindex* index = expr->nodetype() != index_node ? nullptr : static_cast<ASTindex*>(expr);
 
-            default:  break;
-        }
-        if (component != -1) {
-            if (expr->nodetype() == index_node) {
-                static_cast<ASTindex*>(expr)->extend(new ASTliteral (comp, component));
-                return expr;
+        int indexes[3];
+        switch (ASTswizzle::indices (field, indexes, 3, true)) {
+            case 1: {
+                if (!index)
+                    return new ASTindex (comp, expr, new ASTliteral (comp, indexes[0]));
+                index->extend(new ASTliteral (comp, indexes[0]));
+                return index;
             }
-            return new ASTindex (comp, expr, new ASTliteral (comp, component));
+
+            case 3: {
+                // Don't leak soon to be unused expr node
+                std::unique_ptr<ASTNode> cleanup(index);
+                ASTNode* index0 = nullptr;
+                if (index) {
+                    index0 = index->index().get();
+                    expr = index->lvalue().get();
+                }
+
+                ASTNode *args[3];
+                for (int i = 0; i < 3; ++i) {
+                    if (indexes[i] >= 0) {
+                        args[i] = new ASTliteral (comp, indexes[i]);
+                        if (i == 0 && index) {
+                            // Re-use expr by extending the ASTindex.
+                            index->extend (args[i]);
+                            args[0] = cleanup.release ();
+                        } else {
+                            args[i] = !index0 ? new ASTindex (comp, expr, args[i]) :
+                                     new ASTindex (comp, expr, index0, args[i]);
+                         }
+                    } else {
+                        float cval = indexes[i] - ASTswizzle::const_offset;
+                        ASSERT ((cval==0) || (cval==1));
+                        args[i] = new ASTliteral (comp, cval);
+                    }
+                }
+                args[0]->append (args[1]);
+                args[1]->append (args[2]);
+                return new ASTswizzle (comp, args[0], field);
+            }
+
+            default: break;
         }
     }
 

src/liboslcomp/ast.h

@@ -62,7 +62,7 @@ class ASTNode : public OIIO::RefCnt {
         unknown_node, shader_declaration_node, function_declaration_node,
         variable_declaration_node, compound_initializer_node,
         variable_ref_node, preincdec_node, postincdec_node,
-        index_node, structselect_node,
+        index_node, structselect_node, swizzle_node,
         conditional_statement_node,
         loop_statement_node, loopmod_statement_node, return_statement_node,
         binary_expression_node, unary_expression_node,
@@ -644,6 +644,34 @@ class ASTstructselect : public ASTfieldselect
 
 
 
+class ASTswizzle : public ASTfieldselect
+{
+public:
+    ASTswizzle (OSLCompilerImpl *comp, ASTNode *expr, ustring field);
+
+    const char *nodetypename () const { return "swizzle"; }
+    const char *childname (size_t i) const;
+    void print (std::ostream &out, int indentlevel=0) const;
+    TypeSpec typecheck (TypeSpec expected);
+    Symbol *codegen (Symbol *dest = NULL);
+
+    /// Get component indeces for a swizzle string.
+    /// consts allows '0' & '1' to be included in the string.
+    static size_t indices (ustring components, int *indexes, size_t N, bool consts);
+
+    /// Offset if the component index is really a constant ('0' or '1').
+    enum { const_offset = -2 };
+
+    /// Special code generation of assignment of src to proper components.
+    Symbol* codegen_assign (Symbol *src);
+
+    size_t indices (int *indexes, size_t N, bool consts) const {
+        return indices (m_field, indexes, N, consts);
+    }
+};
+
+
+
 class ASTconditional_statement : public ASTNode
 {
 public:

src/liboslcomp/codegen.cpp

@@ -479,6 +479,8 @@ ASTassign_expression::codegen (Symbol *dest)
         // Assigning to an individual component or array element
         index = (ASTindex *) var().get();
         dest = NULL;
+    } else if (var()->nodetype() == swizzle_node) {
+        return static_cast<ASTswizzle*>(var().get())->codegen_assign(expr()->codegen (dest));
     } else if (var()->nodetype() == structselect_node) {
         dest = var()->codegen();
     } else {
@@ -1202,6 +1204,78 @@ ASTindex::codegen_assign (Symbol *src, Symbol *ind,
 
 
 
+Symbol *
+ASTswizzle::codegen (Symbol *dest)
+{
+    // General case, construct a new triple with the swizzle indexes.
+
+    if (dest == nullptr || ! equivalent (dest->typespec(), typespec()))
+        dest = m_compiler->make_temporary (typespec());
+
+    Symbol *syms[4] = { dest };
+
+    int nsym = 1;
+	for (ref arg = child(0); arg && nsym < 4; ++nsym, arg = arg->next()) {
+        syms[nsym] = arg->codegen();
+	}
+	ASSERT (nsym == 4);
+
+    // emit the constructor call
+    emitcode (typespec().string().c_str(), nsym, syms);
+    return dest;
+}
+
+
+
+Symbol *
+ASTswizzle::codegen_assign (Symbol *src)
+{
+    // Swizzle assignment.
+
+    if (!m_is_lvalue) {
+        error ("Cannot assign to constant swizzle");
+        return nullptr;
+    }
+
+    // First child is an index, which holds the lvalue (this) to store to.
+    ASTindex* index = static_cast<ASTindex*>(child(0));
+    Symbol *syms[4] = { index->lvalue().get()->codegen() };
+
+    int nsym = 0;
+    if (src->typespec().is_triple()) {
+        int idxs[3];
+        if (indices(idxs, 3, false /*can't assign to constants*/) != 3) {
+            error ("Trying to assign to invalid swizzle");
+            return nullptr;
+        }
+        // Iterate all of the indexes
+        for (; index && nsym < 3; ++nsym, index = static_cast<ASTindex*>(index->nextptr())) {
+            // tmp[I] = src[N]
+            ASSERT (index->nodetype() == index_node);
+            Symbol* sym = m_compiler->make_temporary (TypeDesc::TypeFloat);
+            emitcode ("compref", sym, src, m_compiler->make_constant (nsym));
+
+            // sym[0] is taken, so assign to syms[i+1]
+            ASSERT (idxs[nsym]+1 < (sizeof(syms)/sizeof(syms[0])));
+            syms[idxs[nsym]+1] = sym;
+        }
+    } else {
+        // Same assignment to all components
+        for (int n = m_typespec.aggregate(); nsym < n; ++nsym)
+            syms[nsym] = src;
+    }
+
+    ASSERT (nsym == 3);
+
+    // emit the constructor call
+    emitcode (typespec().string().c_str(), nsym+1, syms);
+
+    // so transitive assignment will work
+    return syms[0];
+}
+
+
+
 Symbol *
 ASTstructselect::codegen (Symbol *dest)
 {

src/liboslcomp/typecheck.cpp

@@ -321,6 +321,22 @@ ASTindex::typecheck (TypeSpec expected)
 
 
 
+TypeSpec
+ASTswizzle::typecheck (TypeSpec expected)
+{
+    typecheck_children ();
+
+    // Can't be an lvalue if any of the args are constant floats.
+    for (ref arg = child(0); arg; arg = arg->next()) {
+        if (!(m_is_lvalue = arg->is_lvalue()))
+            break;
+    }
+
+    return m_typespec;
+}
+
+
+
 TypeSpec
 ASTstructselect::typecheck (TypeSpec expected)
 {

testsuite/swizzle/ref/out.txt

@@ -1,4 +1,18 @@
 Compiled swizzle.osl -> swizzle.oso
+c: 1 2 3
+c.brg = c.grb = c.bgr: 3 1 2
+c.rgb = c.gbr: 1 2 3
+c.r00: 1 0 0
+c.g01: 2 0 1
+c.b11: 3 1 1
+c.bg1: 3 2 1
+c.rgb = c.brg: 3 1 2
+c.rgb = c.brg: 2 3 1
+c.rgb = c.brg: 1 2 3
+c.rgb = c.bgr: 3 2 1
+c.bgr = c.bgr: 3 2 1
+c.grb = c.bgr: 2 1 3
+c.ggr + c.rgr: 3 2 4
 c.r: 4
 c.g: 3
 c.b: 2
@@ -12,4 +26,15 @@ c4[0].b: 3
 c4[1].r: -1
 c4[1].b: -3
 c4[1].g: -2
+c.rrr: 4 4 4
+c.rgb: 4 3 2
+c.rbg: 4 2 3
+c.ggg: 3 3 3
+c.grb: 3 4 2
+c.gbr: 3 2 4
+c.bbb: 2 2 2
+c.brg: 2 4 3
+c.ggr: 3 3 4
+ca[0].brg: 3 1 2
+c.rgb = 1: 1 1 1
 

testsuite/swizzle/swizzle.osl

@@ -4,7 +4,44 @@
 
 shader swizzle ()
 {
-    color c = color(4,3,2);
+    color c = color(1,2,3);
+    printf("c: %g\n", c);
+
+    // transitive
+    // 1,2,3 -> 2,3,1 -> 3,1,2
+    c.brg = c.grb = c.bgr;
+    printf("c.brg = c.grb = c.bgr: %g\n", c);
+
+    c.rgb = c.gbr;
+    printf("c.rgb = c.gbr: %g\n", c);
+
+    printf("c.r00: %g\n", c.r00);
+    printf("c.g01: %g\n", c.g01);
+    printf("c.b11: %g\n", c.b11);
+    printf("c.bg1: %g\n", c.bg1);
+
+    c = c.brg;
+    printf("c.rgb = c.brg: %g\n", c);
+
+    c = c.brg;
+    printf("c.rgb = c.brg: %g\n", c);
+
+    c = c.brg;
+    printf("c.rgb = c.brg: %g\n", c);
+
+	c.rgb = c.bgr;
+    printf("c.rgb = c.bgr: %g\n", c);
+
+    c.bgr = c.bgr;
+    printf("c.bgr = c.bgr: %g\n", c);
+
+    c.grb = c.bgr;
+    printf("c.grb = c.bgr: %g\n", c);
+
+    printf("c.ggr + c.rgr: %g\n", c.ggr + c.rgr);
+
+    c.gbr = c.ggr + c.rgr;
+
     printf("c.r: %g\n", c.r);
     printf("c.g: %g\n", c.g);
     printf("c.b: %g\n", c.b);
@@ -26,4 +63,19 @@ shader swizzle ()
         printf("c4[%d].%s: %g\n", i, i ? "b" : "g", i ? c4[i].rgb.b : c4[i].rgb.g);
         printf("c4[%d].%s: %g\n", i, i ? "g" : "b", i ? c4[i].rgb.g : c4[i].rgb.b);
     }
+
+    printf("c.rrr: %g\n", c.rrr);
+    printf("c.rgb: %g\n", c.rgb);
+    printf("c.rbg: %g\n", c.rbg);
+    printf("c.ggg: %g\n", c.ggg);
+    printf("c.grb: %g\n", c.grb);
+    printf("c.gbr: %g\n", c.gbr);
+    printf("c.bbb: %g\n", c.bbb);
+    printf("c.brg: %g\n", c.brg);
+    printf("c.ggr: %g\n", c.ggr);
+
+    printf("ca[0].brg: %g\n", ca[0].brg);
+
+    c.rgb = 1;
+    printf("c.rgb = 1: %g\n", c);
 }