From bf60c704da575c98faf0275ae401e0b1019a3141 Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Tue, 2 Jul 2024 14:19:57 -0400
Subject: [PATCH] xo-expression: + ptr to originating lambda

---
 include/xo/expression/LocalEnv.hpp | 22 ++++++++++++++++++++--
 include/xo/expression/Variable.hpp |  5 -----
 src/expression/Lambda.cpp          | 16 ++++++++++++----
 3 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/include/xo/expression/LocalEnv.hpp b/include/xo/expression/LocalEnv.hpp
index b206bddf..2a7a3045 100644
--- a/include/xo/expression/LocalEnv.hpp
+++ b/include/xo/expression/LocalEnv.hpp
@@ -9,6 +9,8 @@
 
 namespace xo {
     namespace ast {
+        class Lambda;
+
         /** @brief LocalEnv
          *
          *  @class Local environment for a lambda.
@@ -26,10 +28,17 @@ namespace xo {
                 return new LocalEnv(argv);
             }
 
+            Lambda * owner() const { return owner_; }
             const std::vector<ref::rp<Variable>> & argv() const { return argv_; }
             int n_arg() const { return argv_.size(); }
             TypeDescr fn_arg(uint32_t i) const { return argv_[i]->valuetype(); }
 
+            /** single-assign this environment's owner **/
+            void assign_owner(Lambda * p) {
+                assert(owner_ == nullptr);
+                owner_ = p;
+            }
+
             /** single-assign this environment's parent **/
             void assign_parent(ref::brw<Environment> p) {
                 assert(parent_env_.get() == nullptr);
@@ -55,11 +64,20 @@ namespace xo {
                 : argv_(argv) {}
 
         private:
+            /** Lambnda for which this environment created.
+             *
+             *  Invariant:
+             *  @code
+             *  owner_->local_env_ == this
+             *  @endcode
+             **/
+            Lambda * owner_ = nullptr;
+
             /** formal argument names **/
             std::vector<ref::rp<Variable>> argv_;
 
-            /** parent environment.  Free variable in this lambda's
-             *  body,  will be resolved by referring them to @ref parent_env_.
+            /** parent environment.  A free variable in this lambda's
+             *  body will be resolved by referring them to @ref parent_env_.
              **/
             ref::rp<Environment> parent_env_;
         };
diff --git a/include/xo/expression/Variable.hpp b/include/xo/expression/Variable.hpp
index e8181a8c..471eb873 100644
--- a/include/xo/expression/Variable.hpp
+++ b/include/xo/expression/Variable.hpp
@@ -44,11 +44,6 @@ namespace xo {
 
             virtual void attach_envs(ref::brw<Environment> /*p*/) override {}
 
-#ifdef NOT_USING
-            virtual std::int32_t find_free_vars(std::set<ref::brw<Variable>> * p_set) override {
-            }
-#endif
-
             virtual void display(std::ostream & os) const override;
 
         private:
diff --git a/src/expression/Lambda.cpp b/src/expression/Lambda.cpp
index d2fe8222..6c8afafd 100644
--- a/src/expression/Lambda.cpp
+++ b/src/expression/Lambda.cpp
@@ -39,10 +39,18 @@ namespace xo {
             TypeDescr lambda_td
                 = TypeDescrBase::require_by_fn_info(function_info);
 
-            return new Lambda(name,
-                              lambda_td,
-                              LocalEnv::make(argv),
-                              body);
+            rp<LocalEnv> env = LocalEnv::make(argv);
+
+            rp<Lambda> retval
+                = new Lambda(name,
+                             lambda_td,
+                             env,
+                             body);
+
+            /* need two-phase construction b/c pointer cycle */
+            env->assign_owner(retval.get());
+
+            return retval;
         } /*make*/
 
         std::set<std::string>