From b7db7c54549d6ad52ecbd1b4ac5feee8b9eeaac3 Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Mon, 1 Jul 2024 16:50:17 -0400
Subject: [PATCH] xo-expression: + Expression::get_free_variables()

---
 include/xo/expression/Apply.hpp              | 14 ++++++++++++++
 include/xo/expression/ConstantInterface.hpp  |  4 ++++
 include/xo/expression/Expression.hpp         |  6 ++++++
 include/xo/expression/IfExpr.hpp             | 15 +++++++++++++++
 include/xo/expression/Lambda.hpp             | 10 ++++++++++
 include/xo/expression/PrimitiveInterface.hpp |  4 ++++
 include/xo/expression/Variable.hpp           |  6 ++++++
 7 files changed, 59 insertions(+)

diff --git a/include/xo/expression/Apply.hpp b/include/xo/expression/Apply.hpp
index 666ca5c7..b6db3a0c 100644
--- a/include/xo/expression/Apply.hpp
+++ b/include/xo/expression/Apply.hpp
@@ -36,6 +36,20 @@ namespace xo {
             const ref::rp<Expression> & fn() const { return fn_; }
             const std::vector<ref::rp<Expression>> & argv() const { return argv_; }
 
+            virtual std::set<std::string> get_free_variables() const override {
+                std::set<std::string> retval = fn_->get_free_variables();
+
+                for (const auto & arg : argv_) {
+                    std::set<std::string> arg_free_set
+                        = arg->get_free_variables();
+
+                    for (const auto & name : arg_free_set)
+                        retval.insert(name);
+                }
+
+                return retval;
+            }
+
             virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
                 std::size_t n = 1;
 
diff --git a/include/xo/expression/ConstantInterface.hpp b/include/xo/expression/ConstantInterface.hpp
index cb38f50d..0b10ec41 100644
--- a/include/xo/expression/ConstantInterface.hpp
+++ b/include/xo/expression/ConstantInterface.hpp
@@ -36,6 +36,10 @@ namespace xo {
 
             // ----- Expression -----
 
+            virtual std::set<std::string> get_free_variables() const override {
+                return std::set<std::string>();
+            }
+
             virtual void attach_envs(ref::brw<Environment> /*p*/) override {}
 
 
diff --git a/include/xo/expression/Expression.hpp b/include/xo/expression/Expression.hpp
index 403c2403..d23a412e 100644
--- a/include/xo/expression/Expression.hpp
+++ b/include/xo/expression/Expression.hpp
@@ -45,6 +45,12 @@ namespace xo {
             exprtype extype() const { return extype_; }
             TypeDescr valuetype() const { return valuetype_; }
 
+            /** find free named variables in this expression.
+             *  comprises the set of names that don't match formal parameters in
+             *  enclosing lambdas.
+             **/
+            virtual std::set<std::string> get_free_variables() const = 0;
+
             /** visit each Expression node in this AST,
              *  and invoke @p fn for each.
              *  Returns the number of nodes visited.
diff --git a/include/xo/expression/IfExpr.hpp b/include/xo/expression/IfExpr.hpp
index f962c905..fc8204a7 100644
--- a/include/xo/expression/IfExpr.hpp
+++ b/include/xo/expression/IfExpr.hpp
@@ -40,6 +40,21 @@ namespace xo {
 
             // ----- Expression -----
 
+            virtual std::set<std::string> get_free_variables() const override {
+                std::set<std::string> retval = test_->get_free_variables();
+
+                std::set<std::string> free_vars;
+                free_vars = when_true_->get_free_variables();
+                for (const auto & s : free_vars)
+                    retval.insert(s);
+
+                free_vars = when_false_->get_free_variables();
+                for (const auto & s : free_vars)
+                    retval.insert(s);
+
+                return retval;
+            }
+
             virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
                 std::size_t n = 1;
 
diff --git a/include/xo/expression/Lambda.hpp b/include/xo/expression/Lambda.hpp
index 791f3548..a7f7529d 100644
--- a/include/xo/expression/Lambda.hpp
+++ b/include/xo/expression/Lambda.hpp
@@ -49,6 +49,16 @@ namespace xo {
 
             // ----- Expression -----
 
+            virtual std::set<std::string> get_free_variables() const override {
+                std::set<std::string> retval = body_->get_free_variables();
+
+                /* but remove formals. */
+                for (const auto & var : local_env_->argv())
+                    retval.erase(var->name());
+
+                return retval;
+            }
+
             virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
                 std::size_t n = 1;
 
diff --git a/include/xo/expression/PrimitiveInterface.hpp b/include/xo/expression/PrimitiveInterface.hpp
index 184a70d8..c52eb979 100644
--- a/include/xo/expression/PrimitiveInterface.hpp
+++ b/include/xo/expression/PrimitiveInterface.hpp
@@ -48,6 +48,10 @@ namespace xo {
 
             // ----- Expression -----
 
+            virtual std::set<std::string> get_free_variables() const override {
+                return std::set<std::string>();
+            }
+
             virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
                 visitor_fn(this);
                 return 1;
diff --git a/include/xo/expression/Variable.hpp b/include/xo/expression/Variable.hpp
index c4a2eb9a..e8181a8c 100644
--- a/include/xo/expression/Variable.hpp
+++ b/include/xo/expression/Variable.hpp
@@ -31,6 +31,12 @@ namespace xo {
 
             const std::string & name() const { return name_; }
 
+            virtual std::set<std::string> get_free_variables() const override {
+                std::set<std::string> retval;
+                retval.insert(this->name_);
+                return retval;
+            }
+
             virtual std::size_t visit_preorder(VisitFn visitor_fn) override {
                 visitor_fn(this);
                 return 1;