- Improved the way iter works (WRT the register offset).

- Fixed some bugs of how the loop-transform was working.

3 files changed, 169 insertions, 60 deletions

diff --git a/src/lux/analyser/base.clj b/src/lux/analyser/base.clj
index 560166eb7..9bdcdeb11 100644
--- a/src/lux/analyser/base.clj
+++ b/src/lux/analyser/base.clj
@@ -86,3 +86,46 @@
 
 (defn |meta [type cursor analysis]
   (&/T [(&/T [type cursor]) analysis]))
+
+(defn de-meta
+  "(-> Analysis Analysis)"
+  [analysis]
+  (|let [[meta analysis-] analysis]
+    (|case analysis-
+      ($variant idx is-last? value)
+      ($variant idx is-last? (de-meta value))
+      
+      ($tuple elems)
+      ($tuple (&/|map de-meta elems))
+      
+      ($apply func args)
+      ($apply (de-meta func)
+              (&/|map de-meta args))
+      
+      ($case value branches)
+      ($case (de-meta value)
+             (&/|map (fn [branch]
+                       (|let [[_pattern _body] branch]
+                         (&/T [_pattern (de-meta _body)])))
+                     branches))
+      
+      ($lambda _register-offset scope captured body)
+      ($lambda _register-offset scope
+          (&/|map (fn [branch]
+                    (|let [[_name _captured] branch]
+                      (&/T [_name (de-meta _captured)])))
+                  captured)
+          (de-meta body))
+
+      ($ann value-expr type-expr)
+      (de-meta value-expr)
+      
+      ($captured scope idx source)
+      ($captured scope idx (de-meta source))
+
+      ($proc proc-ident args special-args)
+      ($proc proc-ident (&/|map de-meta args) special-args)
+      
+      _
+      analysis-
+      )))
diff --git a/src/lux/compiler/lux.clj b/src/lux/compiler/lux.clj
index 74c6221bb..f44375e97 100644
--- a/src/lux/compiler/lux.clj
+++ b/src/lux/compiler/lux.clj
@@ -152,12 +152,10 @@
 
 (defn compile-loop [compile-expression register-offset inits body]
   (|do [^MethodVisitor *writer* &/get-writer
-        :let [_ (.println System/out (pr-str 'compile-loop register-offset (&/|length inits)))]
         :let [idxs+inits (&/zip2 (&/|range* 0 (dec (&/|length inits)))
                                  inits)]
         _ (&/map% (fn [idx+_init]
                     (|do [:let [[idx _init] idx+_init
-                                _ (.println System/out (pr-str 'compile-loop/_init (&/adt->text _init)))
                                 idx+ (+ register-offset idx)]
                           _ (compile-expression nil _init)
                           :let [_ (.visitVarInsn *writer* Opcodes/ASTORE idx+)]]
@@ -170,7 +168,7 @@
 
 (defn compile-iter [compile $begin register-offset ?args]
   (|do [^MethodVisitor *writer* &/get-writer
-        :let [idxs+args (&/zip2 (&/|range* 1 (&/|length ?args))
+        :let [idxs+args (&/zip2 (&/|range* 0 (dec (&/|length ?args)))
                                 ?args)]
         _ (&/map% (fn [idx+?arg]
                     (|do [:let [[idx ?arg] idx+?arg
diff --git a/src/lux/optimizer.clj b/src/lux/optimizer.clj
index 755c4aae2..5c30dc44f 100644
--- a/src/lux/optimizer.clj
+++ b/src/lux/optimizer.clj
@@ -101,6 +101,71 @@
   ;; branch's body.
   ("ExecPM" 1))
 
+(defn de-meta
+  "(-> Optimized Optimized)"
+  [optim]
+  (|let [[meta optim-] optim]
+    (|case optim-
+      ($variant idx is-last? value)
+      ($variant idx is-last? (de-meta value))
+      
+      ($tuple elems)
+      ($tuple (&/|map de-meta elems))
+      
+      ($case value [_pm _bodies])
+      ($case (de-meta value)
+             (&/T [_pm (&/|map de-meta _bodies)]))
+      
+      ($function _register-offset arity scope captured body*)
+      ($function _register-offset
+                 arity
+                 scope
+                 (&/|map (fn [capture]
+                           (|let [[_name [_meta ($captured _scope _idx _source)]] capture]
+                             (&/T [_name ($captured _scope _idx (de-meta _source))])))
+                         captured)
+                 (de-meta body*))
+
+      ($ann value-expr type-expr)
+      (de-meta value-expr)
+      
+      ($apply func args)
+      ($apply (de-meta func)
+              (&/|map de-meta args))
+      
+      ($captured scope idx source)
+      ($captured scope idx (de-meta source))
+      
+      ($proc proc-ident args special-args)
+      ($proc proc-ident (&/|map de-meta args) special-args)
+
+      ($loop _register-offset _inits _body)
+      ($loop _register-offset
+             (&/|map de-meta _inits)
+             (de-meta _body))
+      
+      ($iter _iter-register-offset args)
+      ($iter _iter-register-offset
+             (&/|map de-meta args))
+
+      ($let _value _register _body)
+      ($let (de-meta _value)
+            _register
+            (de-meta _body))
+
+      ($record-get _value _path)
+      ($record-get (de-meta _value)
+                   _path)
+
+      ($if _test _then _else)
+      ($if (de-meta _test)
+           (de-meta _then)
+           (de-meta _else))
+      
+      _
+      optim-
+      )))
+
 ;; This function does a simple transformation from the declarative
 ;; model of PM of the analyser, to the operational model of PM of the
 ;; optimizer.
@@ -797,7 +862,7 @@
     (|case optim-
       ($apply [meta-0 ($var (&/$Local 0))] _args)
       (if (= arity (&/|length _args))
-        (&/T [meta ($iter 0 _args)])
+        (&/T [meta ($iter 1 _args)])
         optim)
 
       ($case _value [_pattern _bodies])
@@ -880,99 +945,102 @@
       false
       )))
 
-(defn ^:private pm-loop-transform [register-offset pattern]
+(defn ^:private pm-loop-transform [register-offset direct? pattern]
   (|case pattern
     ($BindPM _var-id)
-    ($BindPM (+ register-offset _var-id))
+    ($BindPM (+ register-offset (if direct?
+                                  (- _var-id 2)
+                                  (- _var-id 1))))
 
     ($SeqPM _left-pm _right-pm)
-    ($SeqPM (pm-loop-transform register-offset _left-pm)
-            (pm-loop-transform register-offset _right-pm))
+    ($SeqPM (pm-loop-transform register-offset direct? _left-pm)
+            (pm-loop-transform register-offset direct? _right-pm))
 
     ($AltPM _left-pm _right-pm)
-    ($AltPM (pm-loop-transform register-offset _left-pm)
-            (pm-loop-transform register-offset _right-pm))
+    ($AltPM (pm-loop-transform register-offset direct? _left-pm)
+            (pm-loop-transform register-offset direct? _right-pm))
 
     _
     pattern
     ))
 
-(defn ^:private loop-transform [register-offset body]
-  (|let [[meta body-] body]
+;; This function must be run STRICTLY before shift-function body, as
+;; the transformation assumes that SFB will be invoke after it.
+(defn ^:private loop-transform [register-offset direct? body]
+  (|let [adjust-direct (fn [register]
+                         ;; The register must be decreased once, since
+                         ;; it will be re-increased in
+                         ;; shift-function-body.
+                         ;; The decrease is meant to keep things stable.
+                         (if direct?
+                           ;; And, if this adjustment is done
+                           ;; directly during a loop-transform (and
+                           ;; not indirectly if transforming an inner
+                           ;; loop), then it must be decreased again
+                           ;; because the 0/self var will no longer
+                           ;; exist in the loop's context.
+                           (- register 2)
+                           (- register 1)))
+         [meta body-] body]
     (|case body-
       ($variant idx is-last? value)
-      (&/T [meta ($variant idx is-last? (loop-transform register-offset value))])
+      (&/T [meta ($variant idx is-last? (loop-transform register-offset direct? value))])
       
       ($tuple elems)
-      (&/T [meta ($tuple (&/|map (partial loop-transform register-offset) elems))])
+      (&/T [meta ($tuple (&/|map (partial loop-transform register-offset direct?) elems))])
       
       ($case value [_pm _bodies])
-      (&/T [meta ($case (loop-transform register-offset value)
-                        (&/T [(pm-loop-transform register-offset _pm)
-                              (&/|map (partial loop-transform register-offset)
+      (&/T [meta ($case (loop-transform register-offset direct? value)
+                        (&/T [(pm-loop-transform register-offset direct? _pm)
+                              (&/|map (partial loop-transform register-offset direct?)
                                       _bodies)]))])
-      
-      ($function _register-offset arity scope captured body*)
-      (&/T [meta ($function _register-offset
-                            arity
-                            scope
-                            (&/|map (fn [entry]
-                                      (|let [[_name _captured] entry]
-                                        (|case _captured
-                                          ;; [_meta ($captured _scope _idx [_ ($var (&/$Local _l-idx))])]
-                                          ;; (&/T [_meta ($var (&/$Local (+ register-offset (dec _l-idx))))])
-
-                                          [_meta ($captured _scope _idx _source)]
-                                          (&/T [_name _source])
-                                          )))
-                                    captured)
-                            body*)])
 
+      ;; Functions are ignored because they'll be handled properly at shift-function-body
+      
       ($ann value-expr type-expr)
-      (&/T [meta ($ann (loop-transform register-offset value-expr)
+      (&/T [meta ($ann (loop-transform register-offset direct? value-expr)
                        type-expr)])
       
       ($var (&/$Local idx))
-      ;; First, it's decreased because the var index is 1-based (since
-      ;; 0 is reserved for self-reference).
+      ;; The index must be decreased once, because the var index is
+      ;; 1-based (since 0 is reserved for self-reference).
+      ;; Then it must be decreased again, since it will be increased
+      ;; in the shift-function-body call.
       ;; Then, I add the offset to ensure the var points to the right register.
-      (&/T [meta ($var (&/$Local (+ register-offset (dec idx))))])
+      (&/T [meta ($var (&/$Local (-> (adjust-direct idx)
+                                     (+ register-offset))))])
 
       ($apply func args)
-      (&/T [meta ($apply (loop-transform register-offset func)
-                         (&/|map (partial loop-transform register-offset) args))])
+      (&/T [meta ($apply (loop-transform register-offset direct? func)
+                         (&/|map (partial loop-transform register-offset direct?) args))])
       
-      ;; ($captured scope idx [_ ($var (&/$Local _l-idx))])
-      ;; (&/T [meta ($var (&/$Local (+ register-offset (dec _l-idx))))])
-
-      ($captured scope idx source)
-      source
+      ;; Captured-vars are ignored because they'll be handled properly at shift-function-body
       
       ($proc proc-ident args special-args)
-      (&/T [meta ($proc proc-ident (&/|map (partial loop-transform register-offset) args) special-args)])
+      (&/T [meta ($proc proc-ident (&/|map (partial loop-transform register-offset direct?) args) special-args)])
 
       ($loop _register-offset _inits _body)
-      (&/T [meta ($loop (+ register-offset _register-offset)
-                        (&/|map (partial loop-transform register-offset) _inits)
-                        (loop-transform register-offset _body))])
+      (&/T [meta ($loop (+ register-offset (adjust-direct _register-offset))
+                        (&/|map (partial loop-transform register-offset direct?) _inits)
+                        (loop-transform register-offset direct? _body))])
       
       ($iter _iter-register-offset args)
-      (&/T [meta ($iter (+ register-offset _iter-register-offset)
-                        (&/|map (partial loop-transform register-offset) args))])
+      (&/T [meta ($iter (+ register-offset (adjust-direct _iter-register-offset))
+                        (&/|map (partial loop-transform register-offset direct?) args))])
 
       ($let _value _register _body)
-      (&/T [meta ($let (loop-transform register-offset _value)
-                       (+ register-offset _register)
-                       (loop-transform register-offset _body))])
+      (&/T [meta ($let (loop-transform register-offset direct? _value)
+                       (+ register-offset (adjust-direct _register))
+                       (loop-transform register-offset direct? _body))])
 
       ($record-get _value _path)
-      (&/T [meta ($record-get (loop-transform register-offset _value)
+      (&/T [meta ($record-get (loop-transform register-offset direct? _value)
                               _path)])
 
       ($if _test _then _else)
-      (&/T [meta ($if (loop-transform register-offset _test)
-                      (loop-transform register-offset _then)
-                      (loop-transform register-offset _else))])
+      (&/T [meta ($if (loop-transform register-offset direct? _test)
+                      (loop-transform register-offset direct? _then)
+                      (loop-transform register-offset direct? _else))])
       
       _
       body
@@ -980,8 +1048,8 @@
 
 (defn ^:private inline-loop [meta register-offset scope captured args body]
   (->> body
+       (loop-transform register-offset true)
        (shift-function-body scope (&/|tail scope) true)
-       (loop-transform register-offset)
        ($loop register-offset args)
        (list meta)
        (&/T)))
@@ -1033,8 +1101,8 @@
         (|let [=func (pass-0 top-level-func? func)
                =args (&/|map (partial pass-0 top-level-func?) args)]
           (|case =func
-            [_ (&a/$ann [_ ($function _register-offset _arity _scope _captured _body)]
-                        _)]
+            [_ ($ann [_ ($function _register-offset _arity _scope _captured _body)]
+                     _)]
             (if (and (= _arity (&/|length =args))
                      (not (contains-self-reference? _body)))
               (inline-loop meta _register-offset _scope _captured =args _body)