xo-exprssion: + Expression::xform_layer()

include/xo/expression/Apply.hpp

@@ -63,6 +63,15 @@ namespace xo {
                 return n;
             }
 
+            virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
+                this->fn_ = fn_->xform_layer(xform_fn);
+
+                for (auto & arg : argv_)
+                    arg = arg->xform_layer(xform_fn);
+
+                return xform_fn(this);
+            }
+
             virtual void attach_envs(ref::brw<Environment> p) override {
                 fn_->attach_envs(p);
 

include/xo/expression/Constant.hpp

@@ -56,6 +56,10 @@ namespace xo {
                 return 1;
             }
 
+            virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
+                return xform_fn(this);
+            }
+
             virtual void display(std::ostream & os) const override {
                 os << "<Constant"
                    << xtag("type", value_td_->short_name())

include/xo/expression/Expression.hpp

@@ -35,7 +35,10 @@ namespace xo {
          **/
         class Expression : public ref::Refcount {
         public:
-            using VisitFn = std::function<void (ref::brw<Expression>)>;
+            using VisitFn = std::function
+                <void (ref::brw<Expression>)>;
+            using TransformFn = std::function
+                <ref::rp<Expression> (ref::brw<Expression>)>;
             using TypeDescr = xo::reflect::TypeDescr;
 
         public:
@@ -58,6 +61,9 @@ namespace xo {
              **/
             virtual std::size_t visit_preorder(VisitFn visitor_fn) = 0;
 
+            /** traverse ast @ref visit_preorder but do not visit Lambdas **/
+            virtual ref::rp<Expression> xform_layer(TransformFn visitor_fn) = 0;
+
             /** attach an environment to each lambda expression X in this subtree,
              *  that will:
              *  - resolve names matching X's arguments (formal parameters) to

include/xo/expression/IfExpr.hpp

@@ -60,12 +60,21 @@ namespace xo {
 
                 visitor_fn(this);
 
+                n += this->test_->visit_preorder(visitor_fn);
                 n += this->when_true_->visit_preorder(visitor_fn);
                 n += this->when_false_->visit_preorder(visitor_fn);
 
                 return n;
             }
 
+            virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
+                this->test_ = this->test_->xform_layer(xform_fn);
+                this->when_true_ = this->when_true_->xform_layer(xform_fn);
+                this->when_false_= this->when_false_->xform_layer(xform_fn);
+
+                return xform_fn(this);
+            }
+
             virtual void attach_envs(ref::brw<Environment> p) override {
                 test_->attach_envs(p);
                 when_true_->attach_envs(p);

include/xo/expression/Lambda.hpp

@@ -68,6 +68,11 @@ namespace xo {
                 return n;
             }
 
+            virtual ref::rp<Expression> xform_layer(TransformFn /*xform_fn*/) override {
+                /* a layer is bounded by lambdas,  don't enter them */
+                return this;
+            }
+
             virtual void attach_envs(ref::brw<Environment> p) override {
                 local_env_->assign_parent(p);
             }

include/xo/expression/PrimitiveInterface.hpp

@@ -57,6 +57,10 @@ namespace xo {
                 return 1;
             }
 
+            virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
+                return xform_fn(this);
+            }
+
             virtual void attach_envs(ref::brw<Environment> /*p*/) override {}
 
         private:

include/xo/expression/Variable.hpp

@@ -43,6 +43,10 @@ namespace xo {
                 return 1;
             }
 
+            virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
+                return xform_fn(this);
+            }
+
             virtual void attach_envs(ref::brw<Environment> /*p*/) override;
 
             virtual void display(std::ostream & os) const override;