From 829bffd007e7e594115361d87c7c1215b116e3f5 Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Mon, 24 Jun 2024 15:10:06 -0400
Subject: [PATCH] xo-expression: + Expression::visit_preorder()

---
 include/xo/expression/Apply.hpp              | 15 ++++++++++++++-
 include/xo/expression/Constant.hpp           |  5 +++++
 include/xo/expression/Expression.hpp         |  9 +++++++++
 include/xo/expression/IfExpr.hpp             | 11 +++++++++++
 include/xo/expression/Lambda.hpp             | 13 +++++++++++++
 include/xo/expression/PrimitiveInterface.hpp |  7 +++++++
 include/xo/expression/Variable.hpp           |  7 ++++++-
 7 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/include/xo/expression/Apply.hpp b/include/xo/expression/Apply.hpp
index 775c31f9..b7cb8c31 100644
--- a/include/xo/expression/Apply.hpp
+++ b/include/xo/expression/Apply.hpp
@@ -36,7 +36,20 @@ namespace xo {
             const ref::rp<Expression> & fn() const { return fn_; }
             const std::vector<ref::rp<Expression>> & argv() const { return argv_; }
 
-            virtual void display(std::ostream & os) const;
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
+                std::size_t n = 1;
+
+                visitor_fn(this);
+
+                n += fn_->visit_preorder(visitor_fn);
+
+                for (const auto & arg : argv_)
+                    n += arg->visit_preorder(visitor_fn);
+
+                return n;
+            }
+
+            virtual void display(std::ostream & os) const override;
 
         private:
             Apply(TypeDescr apply_valuetype,
diff --git a/include/xo/expression/Constant.hpp b/include/xo/expression/Constant.hpp
index 17d8aa3b..50f4c2ba 100644
--- a/include/xo/expression/Constant.hpp
+++ b/include/xo/expression/Constant.hpp
@@ -51,6 +51,11 @@ namespace xo {
 
             // ----- Expression -----
 
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
+                visitor_fn(this);
+                return 1;
+            }
+
             virtual void display(std::ostream & os) const override {
                 os << "<Constant"
                    << xtag("type", value_td_->short_name())
diff --git a/include/xo/expression/Expression.hpp b/include/xo/expression/Expression.hpp
index 22bd8d4a..b38eecbd 100644
--- a/include/xo/expression/Expression.hpp
+++ b/include/xo/expression/Expression.hpp
@@ -8,6 +8,7 @@
 #include "xo/reflect/TypeDescr.hpp"
 #include "xo/refcnt/Refcounted.hpp"
 #include "exprtype.hpp"
+#include <functional>
 
 namespace xo {
     namespace ast {
@@ -29,6 +30,7 @@ namespace xo {
          **/
         class Expression : public ref::Refcount {
         public:
+            using VisitFn = std::function<void (ref::brw<Expression>)>;
             using TypeDescr = xo::reflect::TypeDescr;
 
         public:
@@ -38,6 +40,13 @@ namespace xo {
             exprtype extype() const { return extype_; }
             TypeDescr valuetype() const { return valuetype_; }
 
+            /** visit each Expression node in this AST,
+             *  and invoke @p fn for each.
+             *  Returns the number of nodes visited.
+             *  Preorder: call @p fn for a node before visiting children
+             **/
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) = 0;
+
             /** write human-readable representation to stream **/
             virtual void display(std::ostream & os) const = 0;
             /** human-readable string representation **/
diff --git a/include/xo/expression/IfExpr.hpp b/include/xo/expression/IfExpr.hpp
index ff9c3e57..6f397b3d 100644
--- a/include/xo/expression/IfExpr.hpp
+++ b/include/xo/expression/IfExpr.hpp
@@ -40,6 +40,17 @@ namespace xo {
 
             // ----- Expression -----
 
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
+                std::size_t n = 1;
+
+                visitor_fn(this);
+
+                n += this->when_true_->visit_preorder(visitor_fn);
+                n += this->when_false_->visit_preorder(visitor_fn);
+
+                return n;
+            }
+
             virtual void display(std::ostream & os) const override;
 
         private:
diff --git a/include/xo/expression/Lambda.hpp b/include/xo/expression/Lambda.hpp
index ce63860d..20c4425b 100644
--- a/include/xo/expression/Lambda.hpp
+++ b/include/xo/expression/Lambda.hpp
@@ -48,6 +48,19 @@ namespace xo {
 
             // ----- Expression -----
 
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
+                std::size_t n = 1;
+
+                visitor_fn(this);
+
+                for (const auto & arg : argv_)
+                    n += arg->visit_preorder(visitor_fn);
+
+                n += body_->visit_preorder(visitor_fn);
+
+                return n;
+            }
+
             virtual void display(std::ostream & os) const override;
 
         private:
diff --git a/include/xo/expression/PrimitiveInterface.hpp b/include/xo/expression/PrimitiveInterface.hpp
index 629772e0..907aef2d 100644
--- a/include/xo/expression/PrimitiveInterface.hpp
+++ b/include/xo/expression/PrimitiveInterface.hpp
@@ -46,6 +46,13 @@ namespace xo {
             // virtual TypeDescr fn_retval() const;
             // virtual TypeDescr fn_arg(uint32_t i) const;
 
+            // ----- Expression -----
+
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
+                visitor_fn(this);
+                return 1;
+            }
+
         private:
         }; /*PrimitiveInterface*/
     } /*namespace ast*/
diff --git a/include/xo/expression/Variable.hpp b/include/xo/expression/Variable.hpp
index 908efbf9..13e46a98 100644
--- a/include/xo/expression/Variable.hpp
+++ b/include/xo/expression/Variable.hpp
@@ -31,7 +31,12 @@ namespace xo {
 
             const std::string & name() const { return name_; }
 
-            virtual void display(std::ostream & os) const;
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
+                visitor_fn(this);
+                return 1;
+            }
+
+            virtual void display(std::ostream & os) const override;
 
         private:
             Variable(const std::string & name,