xo-expression: + Expression::nested_layer()

include/xo/expression/Apply.hpp

@@ -63,6 +63,19 @@ namespace xo {
                 return n;
             }
 
+            virtual std::size_t visit_layer(VisitFn visitor_fn) override {
+                std::size_t n = 1;
+
+                visitor_fn(this);
+
+                n += fn_->visit_layer(visitor_fn);
+
+                for (const auto & arg : argv_)
+                    n += arg->visit_layer(visitor_fn);
+
+                return n;
+            }
+
             virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
                 this->fn_ = fn_->xform_layer(xform_fn);
 

include/xo/expression/Constant.hpp

@@ -56,6 +56,11 @@ namespace xo {
                 return 1;
             }
 
+            virtual std::size_t visit_layer(VisitFn visitor_fn) override {
+                visitor_fn(this);
+                return 1;
+            }
+
             virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
                 return xform_fn(this);
             }

include/xo/expression/Environment.hpp

@@ -6,10 +6,12 @@
 #pragma once
 
 #include "xo/refcnt/Refcounted.hpp"
-#include "Variable.hpp"
+#include "binding_path.hpp"
 
 namespace xo {
     namespace ast {
+        class Expression;
+
         class Environment : public ref::Refcount {
         public:
             /** true if this is toplevel (global) environment.

include/xo/expression/Expression.hpp

@@ -61,6 +61,13 @@ namespace xo {
              **/
             virtual std::size_t visit_preorder(VisitFn visitor_fn) = 0;
 
+            /** visit each Expression node in this AST,
+             *  including immediately-nested Lambda nodes;
+             *  but do not recurse into the params/body of such nested Lambdas.
+             *  Returns the number of nodes visited
+             **/
+            virtual std::size_t visit_layer(VisitFn visitor_fn) = 0;
+
             /** traverse ast @ref visit_preorder but do not visit Lambdas **/
             virtual ref::rp<Expression> xform_layer(TransformFn visitor_fn) = 0;
 

include/xo/expression/IfExpr.hpp

@@ -67,6 +67,18 @@ namespace xo {
                 return n;
             }
 
+            virtual std::size_t visit_layer(VisitFn visitor_fn) override {
+                std::size_t n = 1;
+
+                visitor_fn(this);
+
+                n += this->test_->visit_layer(visitor_fn);
+                n += this->when_true_->visit_layer(visitor_fn);
+                n += this->when_false_->visit_layer(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);

include/xo/expression/Lambda.hpp

@@ -69,6 +69,14 @@ namespace xo {
                 return n;
             }
 
+            virtual std::size_t visit_layer(VisitFn visitor_fn) override {
+                std::size_t n = 1;
+
+                visitor_fn(this);
+
+                return n;
+            }
+
             virtual ref::rp<Expression> xform_layer(TransformFn /*xform_fn*/) override {
                 /* a layer is bounded by lambdas,  don't enter them */
                 return this;

include/xo/expression/LocalEnv.hpp

@@ -6,6 +6,8 @@
 #pragma once
 
 #include "Environment.hpp"
+#include "Variable.hpp"
+#include "xo/reflect/TypeDescr.hpp"
 
 namespace xo {
     namespace ast {

include/xo/expression/PrimitiveInterface.hpp

@@ -57,6 +57,11 @@ namespace xo {
                 return 1;
             }
 
+            virtual std::size_t visit_layer(VisitFn visitor_fn) override {
+                visitor_fn(this);
+                return 1;
+            }
+
             virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
                 return xform_fn(this);
             }

include/xo/expression/Variable.hpp

@@ -48,6 +48,11 @@ namespace xo {
                 return 1;
             }
 
+            virtual std::size_t visit_layer(VisitFn visitor_fn) override {
+                visitor_fn(this);
+                return 1;
+            }
+
             virtual ref::rp<Expression> xform_layer(TransformFn xform_fn) override {
                 return xform_fn(this);
             }

src/expression/Lambda.cpp

@@ -112,7 +112,7 @@ namespace xo {
                                 return var.get();
                             } else {
                                 /* substitute already-encountered var_map[] member */
-                                return ix->second;
+                                return ix->second.get();
                             }
                         } else {
                             return x.get();