From d0b28e3cd4a4ecf15f85a95d720ac978917fa20c Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Fri, 23 Aug 2024 11:35:10 -0400
Subject: [PATCH] xo-expression: + AssignExpr

---
 include/xo/expression/AssignExpr.hpp | 61 ++++++++++++++++++
 include/xo/expression/exprtype.hpp   |  3 +
 src/expression/AssignExpr.cpp        | 93 ++++++++++++++++++++++++++++
 3 files changed, 157 insertions(+)
 create mode 100644 include/xo/expression/AssignExpr.hpp
 create mode 100644 src/expression/AssignExpr.cpp

diff --git a/include/xo/expression/AssignExpr.hpp b/include/xo/expression/AssignExpr.hpp
new file mode 100644
index 00000000..0e7d5e5a
--- /dev/null
+++ b/include/xo/expression/AssignExpr.hpp
@@ -0,0 +1,61 @@
+/* file AssignExpr.hpp
+ *
+ * author: Roland Conybeare, Aug 2024
+ */
+
+#pragma once
+
+#include "Expression.hpp"
+#include "Variable.hpp"
+
+namespace xo {
+    namespace ast {
+        /** @class AssignExpr
+         *  @brief Provide expression for assigning to a variable
+         *
+         *    def pi = 3.14159265;
+         *    def foo = 0.0;
+         *    foo := pi / 2;
+         **/
+        class AssignExpr : public Expression {
+        public:
+            static rp<AssignExpr> make(const rp<Variable> & lhs,
+                                       const rp<Expression> & rhs);
+
+            static ref::brw<AssignExpr> from(ref::brw<Expression> x) {
+                return ref::brw<AssignExpr>::from(x);
+            }
+
+            const rp<Variable> & lhs() const { return lhs_; }
+            const rp<Expression> & rhs() const { return rhs_; }
+
+            std::set<std::string> calc_free_variables() const;
+
+            // ----- inherited from Expression -----
+
+            virtual std::set<std::string> get_free_variables() const override;
+            virtual std::size_t visit_preorder(VisitFn visitor_fn) override;
+            virtual std::size_t visit_layer(VisitFn visitor_fn) override;
+            virtual rp<Expression> xform_layer(TransformFn xform_fn) override;
+            virtual void attach_envs(ref::brw<Environment> p) override;
+
+            virtual void display(std::ostream & os) const override;
+
+        private:
+            AssignExpr(const rp<Variable> & lhs,
+                       const rp<Expression> & rhs);
+
+        private:
+            /** assign to this variable. **/
+            rp<Variable> lhs_;
+            /** assign value of this expression to variable @p lhs **/
+            rp<Expression> rhs_;
+
+            /** free variables for this assignment **/
+            std::set<std::string> free_var_set_;
+        };
+    } /*namespace ast*/
+} /*namespace xo*/
+
+
+/* end AssignExpr.hpp */
diff --git a/include/xo/expression/exprtype.hpp b/include/xo/expression/exprtype.hpp
index ae764021..e9b8393c 100644
--- a/include/xo/expression/exprtype.hpp
+++ b/include/xo/expression/exprtype.hpp
@@ -24,6 +24,8 @@ namespace xo {
             primitive,
             /** variable/function definition **/
             define,
+            /** variable assignment **/
+            assign,
             /** function call **/
             apply,
             /** function definition **/
@@ -49,6 +51,7 @@ namespace xo {
             case exprtype::constant: return "constant";
             case exprtype::primitive: return "primitive";
             case exprtype::define: return "define";
+            case exprtype::assign: return "assign";
             case exprtype::apply: return "apply";
             case exprtype::lambda: return "lambda";
             case exprtype::variable: return "variable";
diff --git a/src/expression/AssignExpr.cpp b/src/expression/AssignExpr.cpp
new file mode 100644
index 00000000..8d56a058
--- /dev/null
+++ b/src/expression/AssignExpr.cpp
@@ -0,0 +1,93 @@
+/* file AssignExpr.cpp
+ *
+ * author: Roland Conybeare
+ */
+
+#include "AssignExpr.hpp"
+#include "xo/indentlog/print/tag.hpp"
+
+namespace xo {
+    namespace ast {
+        rp<AssignExpr>
+        AssignExpr::make(const rp<Variable> & lhs,
+                         const rp<Expression> & rhs)
+        {
+            return new AssignExpr(lhs, rhs);
+        }
+
+        AssignExpr::AssignExpr(const rp<Variable> & lhs,
+                               const rp<Expression> & rhs)
+            : Expression(exprtype::assign, rhs->valuetype()),
+              lhs_{lhs}, rhs_{rhs}
+        {
+            this->free_var_set_ = this->calc_free_variables();
+        }
+
+        std::set<std::string>
+        AssignExpr::calc_free_variables() const
+        {
+            std::set<std::string> retval = lhs_->get_free_variables();
+
+            std::set<std::string> tmp = rhs_->get_free_variables();
+
+            for (const auto & name : tmp)
+                retval.insert(name);
+
+            return retval;
+        }
+
+        std::set<std::string>
+        AssignExpr::get_free_variables() const {
+            return free_var_set_;
+        }
+
+        std::size_t
+        AssignExpr::visit_preorder(VisitFn visitor_fn) {
+            std::size_t n = 1;
+
+            visitor_fn(this);
+
+            n += lhs_->visit_preorder(visitor_fn);
+            n += rhs_->visit_preorder(visitor_fn);
+
+            return n;
+        }
+
+        std::size_t
+        AssignExpr::visit_layer(VisitFn visitor_fn) {
+            std::size_t n = 1;
+
+            visitor_fn(this);
+
+            n += lhs_->visit_layer(visitor_fn);
+            n += rhs_->visit_layer(visitor_fn);
+
+            return n;
+        }
+
+        rp<Expression>
+        AssignExpr::xform_layer(TransformFn xform_fn) {
+            this->lhs_ = Variable::from(lhs_->xform_layer(xform_fn)).promote();
+            this->rhs_ = rhs_->xform_layer(xform_fn);
+
+            return xform_fn(this);
+        }
+
+         void
+         AssignExpr::attach_envs(ref::brw<Environment> p) {
+             lhs_->attach_envs(p);
+             rhs_->attach_envs(p);
+         }
+
+        void
+        AssignExpr::display(std::ostream & os) const {
+            os << "<Assign"
+               << xtag("lhs", lhs_)
+               << xtag("rhs", rhs_)
+               << ">";
+        }
+    } /*namespace ast*/
+} /*namespace xo*/
+
+
+/* end AssignExpr.cpp */