Be smarter about capturing the environment for memoized code

* libguile/memoize.h (SCM_M_CAPTURE_MODULE)
* libguile/memoize.c (MAKMEMO_CAPTURE_MODULE, capture_env):
  (maybe_makmemo_capture_module, memoize): Determine when to capture the
  module on the environment chain at compile-time, instead of at
  runtime.  Introduces a new memoized expression type, capture-module.
  (scm_memoized_expression): Start memoizing with #f as the
  environment.
  (unmemoize): Add unmemoizer.
  (scm_memoize_variable_access_x): Cope with #f as module, and treat as
  the root module (captured before modules were booted).

* libguile/eval.c (eval):
* module/ice-9/eval.scm (primitive-eval): Adapt.

4 files changed, 72 insertions, 61 deletions

diff --git a/libguile/eval.c b/libguile/eval.c
index 43a182a5a..1572c8755 100644
--- a/libguile/eval.c
+++ b/libguile/eval.c
@@ -245,18 +245,6 @@ truncate_values (SCM x)
 }
 #define EVAL1(x, env) (truncate_values (eval ((x), (env))))
 
-/* the environment:
-   (VAL ... . MOD)
-   If MOD is #f, it means the environment was captured before modules were
-   booted.
-   If MOD is the literal value '(), we are evaluating at the top level, and so
-   should track changes to the current module. You have to be careful in this
-   case, because further lexical contours should capture the current module.
-*/
-#define CAPTURE_ENV(env)                                        \
-  (scm_is_null (env) ? scm_current_module () :                  \
-   (scm_is_false (env) ? scm_the_root_module () : env))
-
 static SCM
 eval (SCM x, SCM env)
 {
@@ -288,8 +276,7 @@ eval (SCM x, SCM env)
         SCM new_env;
         int i;
 
-        new_env = make_env (VECTOR_LENGTH (inits), SCM_UNDEFINED,
-                            CAPTURE_ENV (env));
+        new_env = make_env (VECTOR_LENGTH (inits), SCM_UNDEFINED, env);
         for (i = 0; i < VECTOR_LENGTH (inits); i++)
           env_set (new_env, 0, i, EVAL1 (VECTOR_REF (inits, i), env));
         env = new_env;
@@ -298,7 +285,7 @@ eval (SCM x, SCM env)
       }
           
     case SCM_M_LAMBDA:
-      RETURN_BOOT_CLOSURE (mx, CAPTURE_ENV (env));
+      RETURN_BOOT_CLOSURE (mx, env);
 
     case SCM_M_QUOTE:
       return mx;
@@ -307,6 +294,9 @@ eval (SCM x, SCM env)
       scm_define (CAR (mx), EVAL1 (CDR (mx), env));
       return SCM_UNSPECIFIED;
 
+    case SCM_M_CAPTURE_MODULE:
+      return eval (mx, scm_current_module ());
+
     case SCM_M_APPLY:
       /* Evaluate the procedure to be applied.  */
       proc = EVAL1 (CAR (mx), env);
@@ -405,8 +395,7 @@ eval (SCM x, SCM env)
       else
         {
           env = env_tail (env);
-          return SCM_VARIABLE_REF
-            (scm_memoize_variable_access_x (x, CAPTURE_ENV (env)));
+          return SCM_VARIABLE_REF (scm_memoize_variable_access_x (x, env));
         }
 
     case SCM_M_TOPLEVEL_SET:
@@ -421,9 +410,7 @@ eval (SCM x, SCM env)
         else
           {
             env = env_tail (env);
-            SCM_VARIABLE_SET
-              (scm_memoize_variable_access_x (x, CAPTURE_ENV (env)),
-               val);
+            SCM_VARIABLE_SET (scm_memoize_variable_access_x (x, env), val);
             return SCM_UNSPECIFIED;
           }
       }
@@ -654,7 +641,7 @@ scm_c_primitive_eval (SCM exp)
 {
   if (!SCM_EXPANDED_P (exp))
     exp = scm_call_1 (scm_current_module_transformer (), exp);
-  return eval (scm_memoize_expression (exp), SCM_EOL);
+  return eval (scm_memoize_expression (exp), SCM_BOOL_F);
 }
 
 static SCM var_primitive_eval;
diff --git a/libguile/memoize.c b/libguile/memoize.c
index 6eb36d437..5c7129feb 100644
--- a/libguile/memoize.c
+++ b/libguile/memoize.c
@@ -131,6 +131,8 @@ scm_t_bits scm_tc16_memoized;
   MAKMEMO (SCM_M_QUOTE, exp)
 #define MAKMEMO_DEFINE(var, val) \
   MAKMEMO (SCM_M_DEFINE, scm_cons (var, val))
+#define MAKMEMO_CAPTURE_MODULE(exp) \
+  MAKMEMO (SCM_M_CAPTURE_MODULE, exp)
 #define MAKMEMO_APPLY(proc, args)\
   MAKMEMO (SCM_M_APPLY, scm_list_2 (proc, args))
 #define MAKMEMO_CONT(proc) \
@@ -166,6 +168,7 @@ static const char *const memoized_tags[] =
   "let",
   "quote",
   "define",
+  "capture-module",
   "apply",
   "call/cc",
   "call-with-values",
@@ -240,6 +243,22 @@ memoize_exps (SCM exps, SCM env)
 }
   
 static SCM
+capture_env (SCM env)
+{
+  if (scm_is_false (env))
+    return SCM_BOOL_T;
+  return env;
+}
+
+static SCM
+maybe_makmemo_capture_module (SCM exp, SCM env)
+{
+  if (scm_is_false (env))
+    return MAKMEMO_CAPTURE_MODULE (exp);
+  return exp;
+}
+
+static SCM
 memoize (SCM exp, SCM env)
 {
   if (!SCM_EXPANDED_P (exp))
@@ -255,7 +274,9 @@ memoize (SCM exp, SCM env)
 
     case SCM_EXPANDED_PRIMITIVE_REF:
       if (scm_is_eq (scm_current_module (), scm_the_root_module ()))
-        return MAKMEMO_TOP_REF (REF (exp, PRIMITIVE_REF, NAME));
+        return maybe_makmemo_capture_module
+          (MAKMEMO_TOP_REF (REF (exp, PRIMITIVE_REF, NAME)),
+           env);
       else
         return MAKMEMO_MOD_REF (list_of_guile, REF (exp, PRIMITIVE_REF, NAME),
                                 SCM_BOOL_F);
@@ -279,11 +300,15 @@ memoize (SCM exp, SCM env)
                               REF (exp, MODULE_SET, PUBLIC));
 
     case SCM_EXPANDED_TOPLEVEL_REF:
-      return MAKMEMO_TOP_REF (REF (exp, TOPLEVEL_REF, NAME));
+      return maybe_makmemo_capture_module
+        (MAKMEMO_TOP_REF (REF (exp, TOPLEVEL_REF, NAME)), env);
 
     case SCM_EXPANDED_TOPLEVEL_SET:
-      return MAKMEMO_TOP_SET (REF (exp, TOPLEVEL_SET, NAME),
-                              memoize (REF (exp, TOPLEVEL_SET, EXP), env));
+      return maybe_makmemo_capture_module
+        (MAKMEMO_TOP_SET (REF (exp, TOPLEVEL_SET, NAME),
+                          memoize (REF (exp, TOPLEVEL_SET, EXP),
+                                   capture_env (env))),
+         env);
 
     case SCM_EXPANDED_TOPLEVEL_DEFINE:
       return MAKMEMO_DEFINE (REF (exp, TOPLEVEL_DEFINE, NAME),
@@ -343,7 +368,9 @@ memoize (SCM exp, SCM env)
                  && scm_is_eq (name, scm_from_latin1_symbol ("pop-fluid")))
           return MAKMEMO_CALL (MAKMEMO_QUOTE (pop_fluid), 0, SCM_EOL);
         else if (scm_is_eq (scm_current_module (), scm_the_root_module ()))
-          return MAKMEMO_CALL (MAKMEMO_TOP_REF (name), nargs, args);
+          return MAKMEMO_CALL (maybe_makmemo_capture_module
+                               (MAKMEMO_TOP_REF (name), env),
+                               nargs, args);
         else
           return MAKMEMO_CALL (MAKMEMO_MOD_REF (list_of_guile, name,
                                                 SCM_BOOL_F),
@@ -381,11 +408,11 @@ memoize (SCM exp, SCM env)
              meta);
         else
           {
-            proc = memoize (body, env);
+            proc = memoize (body, capture_env (env));
             SCM_SETCAR (SCM_CDR (SCM_MEMOIZED_ARGS (proc)), meta);
           }
 
-	return proc;
+	return maybe_makmemo_capture_module (proc, env);
       }
 
     case SCM_EXPANDED_LAMBDA_CASE:
@@ -462,11 +489,12 @@ memoize (SCM exp, SCM env)
         varsv = scm_vector (vars);
         inits = scm_c_make_vector (VECTOR_LENGTH (varsv),
                                    SCM_BOOL_F);
-        new_env = scm_cons (varsv, env);
+        new_env = scm_cons (varsv, capture_env (env));
         for (i = 0; scm_is_pair (exps); exps = CDR (exps), i++)
           VECTOR_SET (inits, i, memoize (CAR (exps), env));
 
-        return MAKMEMO_LET (inits, memoize (body, new_env));
+        return maybe_makmemo_capture_module
+          (MAKMEMO_LET (inits, memoize (body, new_env)), env);
       }
 
     case SCM_EXPANDED_LETREC:
@@ -484,7 +512,7 @@ memoize (SCM exp, SCM env)
         expsv = scm_vector (exps);
 
         undefs = scm_c_make_vector (nvars, MAKMEMO_QUOTE (SCM_UNDEFINED));
-        new_env = scm_cons (varsv, env);
+        new_env = scm_cons (varsv, capture_env (env));
 
         if (in_order_p)
           {
@@ -495,7 +523,8 @@ memoize (SCM exp, SCM env)
                 body_exps = MAKMEMO_SEQ (MAKMEMO_LEX_SET (make_pos (0, i), init),
                                          body_exps);
               }
-            return MAKMEMO_LET (undefs, body_exps);
+            return maybe_makmemo_capture_module
+              (MAKMEMO_LET (undefs, body_exps), env);
           }
         else
           {
@@ -518,9 +547,11 @@ memoize (SCM exp, SCM env)
             if (scm_is_false (sets))
               return memoize (body, env);
 
-            return MAKMEMO_LET (undefs,
-                                MAKMEMO_SEQ (MAKMEMO_LET (inits, sets),
-                                             memoize (body, new_env)));
+            return maybe_makmemo_capture_module
+              (MAKMEMO_LET (undefs,
+                            MAKMEMO_SEQ (MAKMEMO_LET (inits, sets),
+                                         memoize (body, new_env))),
+               env);
           }
       }
 
@@ -538,7 +569,7 @@ SCM_DEFINE (scm_memoize_expression, "memoize-expression", 1, 0, 0,
 #define FUNC_NAME s_scm_memoize_expression
 {
   SCM_ASSERT_TYPE (SCM_EXPANDED_P (exp), exp, 1, FUNC_NAME, "expanded");
-  return memoize (exp, scm_current_module ());
+  return memoize (exp, SCM_BOOL_F);
 }
 #undef FUNC_NAME
 
@@ -612,6 +643,9 @@ unmemoize (const SCM expr)
                          unmemoize (CAR (args)), unmemoize (CDR (args)));
     case SCM_M_DEFINE:
       return scm_list_3 (scm_sym_define, CAR (args), unmemoize (CDR (args)));
+    case SCM_M_CAPTURE_MODULE:
+      return scm_list_2 (scm_from_latin1_symbol ("capture-module"),
+                         unmemoize (args));
     case SCM_M_IF:
       return scm_list_4 (scm_sym_if, unmemoize (scm_car (args)),
                          unmemoize (scm_cadr (args)), unmemoize (scm_cddr (args)));
@@ -735,6 +769,9 @@ SCM_DEFINE (scm_memoize_variable_access_x, "memoize-variable-access!", 2, 0, 0,
 {
   SCM mx = SCM_MEMOIZED_ARGS (m);
 
+  if (scm_is_false (mod))
+    mod = scm_the_root_module ();
+
   switch (SCM_MEMOIZED_TAG (m))
     {
     case SCM_M_TOPLEVEL_REF:
diff --git a/libguile/memoize.h b/libguile/memoize.h
index 95e92a3a9..68dcd2167 100644
--- a/libguile/memoize.h
+++ b/libguile/memoize.h
@@ -69,6 +69,7 @@ enum
     SCM_M_LET,
     SCM_M_QUOTE,
     SCM_M_DEFINE,
+    SCM_M_CAPTURE_MODULE,
     SCM_M_APPLY,
     SCM_M_CONT,
     SCM_M_CALL_WITH_VALUES,
diff --git a/module/ice-9/eval.scm b/module/ice-9/eval.scm
index ed5103955..e34c08715 100644
--- a/module/ice-9/eval.scm
+++ b/module/ice-9/eval.scm
@@ -43,20 +43,6 @@
 
 
 (eval-when (compile)
-  (define-syntax capture-env
-    (syntax-rules ()
-      ((_ (exp ...))
-       (let ((env (exp ...)))
-         (capture-env env)))
-      ((_ env)
-       (if (null? env)
-           (current-module)
-           (if (not env)
-               ;; the and current-module checks that modules are booted,
-               ;; and thus the-root-module is defined
-               (and (current-module) the-root-module)
-               env)))))
-
   (define-syntax env-toplevel
     (syntax-rules ()
       ((_ env)
@@ -459,8 +445,7 @@
          (variable-ref
           (if (variable? var-or-sym)
               var-or-sym
-              (memoize-variable-access! exp
-                                        (capture-env (env-toplevel env))))))
+              (memoize-variable-access! exp (env-toplevel env)))))
 
         (('if (test consequent . alternate))
          (if (eval test env)
@@ -472,7 +457,7 @@
 
         (('let (inits . body))
          (let* ((width (vector-length inits))
-                (new-env (make-env width #f (capture-env env))))
+                (new-env (make-env width #f env)))
            (let lp ((i 0))
              (when (< i width)
                (env-set! new-env 0 i (eval (vector-ref inits i) env))
@@ -482,11 +467,10 @@
         (('lambda (body meta nreq . tail))
          (let ((proc
                 (if (null? tail)
-                    (make-fixed-closure eval nreq body (capture-env env))
+                    (make-fixed-closure eval nreq body env)
                     (if (null? (cdr tail))
-                        (make-rest-closure eval nreq body (capture-env env))
-                        (apply make-general-closure (capture-env env)
-                               body nreq tail)))))
+                        (make-rest-closure eval nreq body env)
+                        (apply make-general-closure env body nreq tail)))))
            (let lp ((meta meta))
              (unless (null? meta)
                (set-procedure-property! proc (caar meta) (cdar meta))
@@ -518,13 +502,15 @@
          (begin
            (define! name (eval x env))
            (if #f #f)))
-      
+
+        (('capture-module x)
+         (eval x (current-module)))
+
         (('toplevel-set! (var-or-sym . x))
          (variable-set!
           (if (variable? var-or-sym)
               var-or-sym
-              (memoize-variable-access! exp
-                                        (capture-env (env-toplevel env))))
+              (memoize-variable-access! exp (env-toplevel env)))
           (eval x env)))
       
         (('call-with-prompt (tag thunk . handler))
@@ -551,4 +537,4 @@
         (if (macroexpanded? exp)
             exp
             ((module-transformer (current-module)) exp)))
-       '()))))
+       #f))))