From 073bb617eb7e5f76269ca6dba0fe498baff6f058 Mon Sep 17 00:00:00 2001
From: Andy Wingo <wingo@pobox.com>
Date: Thu, 14 May 2009 00:11:25 +0200
Subject: [PATCH] add lexical analyzer and allocator

* module/language/tree-il/optimize.scm: Rework to just export the
  optimize! procedure.

* module/language/tree-il/compile-glil.scm (analyze-lexicals): New
  function, analyzes and allocates lexical variables. Almost ready to
  compile now.
  (codegen): Dedent.
---
 module/language/tree-il/compile-glil.scm | 627 +++++++++++++++--------
 module/language/tree-il/optimize.scm     |   9 +-
 2 files changed, 415 insertions(+), 221 deletions(-)

diff --git a/module/language/tree-il/compile-glil.scm b/module/language/tree-il/compile-glil.scm
index d75ae7a56..f54da31f0 100644
--- a/module/language/tree-il/compile-glil.scm
+++ b/module/language/tree-il/compile-glil.scm
@@ -23,13 +23,196 @@
   #:use-module (system base syntax)
   #:use-module (language glil)
   #:use-module (language tree-il)
+  #:use-module (language tree-il optimize)
   #:use-module (ice-9 common-list)
   #:export (compile-glil))
 
+;; parents: lambda -> parent
+;;  useful when we see a closed-over var, so we can calculate its
+;;  coordinates (depth and index).
+;; bindings: lambda -> (sym ...)
+;;  useful for two reasons: one, so we know how much space to allocate
+;;  when we go into a lambda; and two, so that we know when to stop,
+;;  when looking for closed-over vars.
+;; heaps: sym -> lambda
+;;  allows us to heapify vars in an O(1) fashion
+
+;; allocation: the process of assigning a type and index to each var
+;; a var is external if it is heaps; assigning index is easy
+;; args are assigned in order
+;; locals are indexed as their linear position in the binding path
+;; (let (0 1)
+;;   (let (2 3) ...)
+;;   (let (2) ...))
+;;   (let (2 3 4) ...))
+;; etc.
+
+;; allocation:
+;;  sym -> (local . index) | (heap level . index)
+
+
+(define (analyze-lexicals x)
+  (define (find-diff parent this)
+    (let lp ((parent parent) (n 0))
+      (if (eq? parent this)
+          n
+          (lp (hashq-ref parents parent) (1+ n)))))
+
+  (define (find-heap sym parent)
+    ;; fixme: check displaced lexicals here?
+    (if (memq sym (hashq-ref bindings parent))
+        parent
+        (find-binder sym (hashq-ref parents parent))))
+
+  (define (analyze! x parent level)
+    (define (step y) (analyze! y parent level))
+    (define (recur x parent) (analyze! x parent (1+ level)))
+    (record-case x
+      ((<application> proc args)
+       (step proc) (for-each step args))
+
+      ((<conditional> test then else)
+       (step test) (step then) (step else))
+
+      ((<lexical-ref> name gensym)
+       (if (and (not (memq gensym (hashq-ref bindings parent)))
+                (not (hashq-ref heaps gensym)))
+           (hashq-set! heaps gensym (find-heap gensym parent level))))
+      
+      ((<lexical-set> name gensym exp)
+       (step exp)
+       (if (not (hashq-ref heaps gensym))
+           (hashq-set! heaps gensym (find-heap gensym parent level))))
+      
+      ((<module-set> mod name public? exp)
+       (step exp))
+      
+      ((<toplevel-set> name exp)
+       (step exp))
+      
+      ((<toplevel-define> name exp)
+       (step exp))
+      
+      ((<sequence> exps)
+       (for-each step exps))
+      
+      ((<lambda> vars meta body)
+       (hashq-set! parents x parent)
+       (hashq-set! bindings x
+                   (let rev* ((vars vars) (out '()))
+                     (cond ((null? vars) out)
+                           ((pair? vars) (rev* (cdr vars)
+                                               (cons (car vars) out)))
+                           (else (cons vars out)))))
+       (recur body x)
+       (hashq-set! bindings x (reverse! (hashq-ref bindings x))))
+
+      ((<let> vars vals exp)
+       (for-each step vals)
+       (hashq-set! bindings parent
+                   (append (reverse vars) (hashq-ref bindings parent)))
+       (step exp))
+      
+      ((<letrec> vars vals exp)
+       (hashq-set! bindings parent
+                   (append (reverse vars) (hashq-ref bindings parent)))
+       (for-each step vals)
+       (step exp))
+
+      (else #f)))
+
+    (define (allocate-heap! binder)
+      (hashq-set! heap-indexes binder
+                  (1+ (hashq-ref heap-indexes binder -1))))
+
+    (define (allocate! x level n)
+      (define (step y) (allocate! y level n))
+      (record-case x
+        ((<application> proc args)
+         (step proc) (for-each step args))
+
+        ((<conditional> test then else)
+         (step test) (step then) (step else))
+
+        ((<lexical-set> name gensym exp)
+         (step exp))
+        
+        ((<module-set> mod name public? exp)
+         (step exp))
+        
+        ((<toplevel-set> name exp)
+         (step exp))
+        
+        ((<toplevel-define> name exp)
+         (step exp))
+        
+        ((<sequence> exps)
+         (for-each step exps))
+        
+        ((<lambda> vars meta body)
+         (let lp ((vars vars) (n 0))
+           (if (null? vars)
+               (allocate! body (1+ level) n)
+               (let ((v (if (pair? vars) (car vars) vars)))
+                 (let ((binder (hashq-ref heaps v)))
+                   (hashq-set!
+                    allocation v
+                    (if binder
+                        (cons* 'heap (1+ level) (allocate-heap! binder))
+                        (cons 'stack n))))
+                 (lp (if (pair? vars) (cdr vars) '()) (1+ n))))))
+
+        ((<let> vars vals exp)
+         (for-each step vals)
+         (let lp ((vars vars) (n n))
+           (if (null? vars)
+               (allocate! exp level n)
+               (let ((v (car vars)))
+                 (let ((binder (hashq-ref heaps v)))
+                   (hashq-set!
+                    allocation v
+                    (if binder
+                        (cons* 'heap level (allocate-heap! binder))
+                        (cons 'stack n))))
+                 (lp (cdr vars) (1+ n))))))
+        
+        ((<letrec> vars vals exp)
+         (let lp ((vars vars) (n n))
+           (if (null? vars)
+               (begin
+                 (for-each (lambda (x) (allocate! x level n))
+                           vals)
+                 (allocate! exp level n))
+               (let ((v (car vars)))
+                 (let ((binder (hashq-ref heaps v)))
+                   (hashq-set!
+                    allocation v
+                    (if binder
+                        (cons* 'heap level (allocate-heap! binder))
+                        (cons 'stack n))))
+                 (lp (cdr vars) (1+ n))))))
+
+        (else #f)))
+
+  (define parents (make-hash-table))
+  (define bindings (make-hash-table))
+  (define heaps (make-hash-table))
+  (define allocation (make-hash-table))
+  (define heap-indexes (make-hash-table))
+
+  (hashq-set! bindings #f '())
+  (analyze! x #f 0)
+  (allocate! x 0 0)
+
+  allocation)
+
 (define (compile-glil x e opts)
-  (values (codegen x)
-          (and e (cons (car e) (cddr e)))
-          e))
+  (let ((x (optimize! x e opts)))
+    (let ((allocation (analyze-lexicals x)))
+      (values (codegen (make-lambda (tree-il-src x) '() '() x)
+                       allocation)
+              (and e (cons (car e) (cddr e)))
+              e))))
 
 
 
@@ -57,226 +240,230 @@
                        (eq? (ghil-var-kind var) 'public)))
     (else (error "Unknown kind of variable:" var))))
 
-(define (codegen ghil)
-  (let ((stack '()))
-    (define (push-code! src code)
-      (set! stack (cons code stack))
-      (if src (set! stack (cons (make-glil-source src) stack))))
-    (define (var->binding var)
-      (list (ghil-var-name var) (ghil-var-kind var) (ghil-var-index var)))
-    (define (push-bindings! src vars)
-      (if (not (null? vars))
-          (push-code! src (make-glil-bind (map var->binding vars)))))
-    (define (comp tree tail drop)
-      (define (push-label! label)
-	(push-code! #f (make-glil-label label)))
-      (define (push-branch! src inst label)
-	(push-code! src (make-glil-branch inst label)))
-      (define (push-call! src inst args)
-	(for-each comp-push args)
-	(push-code! src (make-glil-call inst (length args))))
-      ;; possible tail position
-      (define (comp-tail tree) (comp tree tail drop))
-      ;; push the result
-      (define (comp-push tree) (comp tree #f #f))
-      ;; drop the result
-      (define (comp-drop tree) (comp tree #f #t))
-      ;; drop the result if unnecessary
-      (define (maybe-drop)
-	(if drop (push-code! #f *ia-drop*)))
-      ;; return here if necessary
-      (define (maybe-return)
-	(if tail (push-code! #f *ia-return*)))
-      ;; return this code if necessary
-      (define (return-code! src code)
-	(if (not drop) (push-code! src code))
-	(maybe-return))
-      ;; return void if necessary
-      (define (return-void!)
-	(return-code! #f *ia-void*))
-      ;; return object if necessary
-      (define (return-object! src obj)
-	(return-code! src (make-glil-const obj)))
-      ;;
-      ;; dispatch
-      (record-case tree
-	((<ghil-void>)
-	 (return-void!))
-
-	((<ghil-quote> env src obj)
-	 (return-object! src obj))
-
-	((<ghil-ref> env src var)
-	 (return-code! src (make-glil-var 'ref env var)))
-
-	((<ghil-set> env src var val)
-	 (comp-push val)
-	 (push-code! src (make-glil-var 'set env var))
-	 (return-void!))
-
-	((<toplevel-define> src name exp)
-	 (comp-push exp)
-	 (push-code! src (make-glil-var 'define env var))
-	 (return-void!))
-
-	((<conditional> src test then else)
-	 ;;     TEST
-	 ;;     (br-if-not L1)
-	 ;;     THEN
-	 ;;     (br L2)
-	 ;; L1: ELSE
-	 ;; L2:
-	 (let ((L1 (make-label)) (L2 (make-label)))
-	   (comp-push test)
-	   (push-branch! src 'br-if-not L1)
-	   (comp-tail then)
-	   (if (not tail) (push-branch! #f 'br L2))
-	   (push-label! L1)
-	   (comp-tail else)
-	   (if (not tail) (push-label! L2))))
-
-	((<sequence> src exps)
-	 ;; EXPS...
-	 ;; TAIL
-	 (if (null? exps)
-	     (return-void!)
-	     (do ((exps exps (cdr exps)))
-		 ((null? (cdr exps))
-		  (comp-tail (car exps)))
-	       (comp-drop (car exps)))))
-
-	((<let> src vars vals body)
-	 ;; VALS...
-	 ;; (set VARS)...
-	 ;; BODY
-	 (for-each comp-push vals)
-         (push-bindings! src vars)
-	 (for-each (lambda (var) (push-code! #f (make-glil-var 'set env var)))
-		   (reverse vars))
-	 (comp-tail body)
-	 (push-code! #f (make-glil-unbind)))
-
-	((<ghil-mv-bind> env src producer vars rest body)
-	 ;; VALS...
-	 ;; (set VARS)...
-	 ;; BODY
-         (let ((MV (make-label)))
-           (comp-push producer)
-           (push-code! src (make-glil-mv-call 0 MV))
-           (push-code! #f (make-glil-const 1))
-           (push-label! MV)
-           (push-code! #f (make-glil-mv-bind (map var->binding vars) rest))
-           (for-each (lambda (var) (push-code! #f (make-glil-var 'set env var)))
-                     (reverse vars)))
-         (comp-tail body)
-         (push-code! #f (make-glil-unbind)))
-
-	((<ghil-lambda> env src vars rest meta body)
-	 (return-code! src (codegen tree)))
-
-	((<ghil-inline> env src inline args)
-	 ;; ARGS...
-	 ;; (INST NARGS)
-         (let ((tail-table '((call . goto/args)
-                             (apply . goto/apply)
-                             (call/cc . goto/cc))))
-           (cond ((and tail (assq-ref tail-table inline))
-                  => (lambda (tail-inst)
-                       (push-call! src tail-inst args)))
-                 (else
-                  (push-call! src inline args)
-                  (maybe-drop)
-                  (maybe-return)))))
-
-        ((<ghil-values> env src values)
-         (cond (tail ;; (lambda () (values 1 2))
-                (push-call! src 'return/values values))
-               (drop ;; (lambda () (values 1 2) 3)
-                (for-each comp-drop values))
-               (else ;; (lambda () (list (values 10 12) 1))
-                (push-code! #f (make-glil-const 'values))
-                (push-code! #f (make-glil-call 'link-now 1))
-                (push-code! #f (make-glil-call 'variable-ref 0))
-                (push-call! src 'call values))))
-                
-        ((<ghil-values*> env src values)
-         (cond (tail ;; (lambda () (apply values '(1 2)))
-                (push-call! src 'return/values* values))
-               (drop ;; (lambda () (apply values '(1 2)) 3)
-                (for-each comp-drop values))
-               (else ;; (lambda () (list (apply values '(10 12)) 1))
-                (push-code! #f (make-glil-const 'values))
-                (push-code! #f (make-glil-call 'link-now 1))
-                (push-code! #f (make-glil-call 'variable-ref 0))
-                (push-call! src 'apply values))))
-                
-	((<ghil-call> env src proc args)
-	 ;; PROC
-	 ;; ARGS...
-	 ;; ([tail-]call NARGS)
-	 (comp-push proc)
-         (let ((nargs (length args)))
-           (cond ((< nargs 255)
-                  (push-call! src (if tail 'goto/args 'call) args))
-                 (else
-                  (push-call! src 'mark '())
-                  (for-each comp-push args)
-                  (push-call! src 'list-mark '())
-                  (push-code! src (make-glil-call (if tail 'goto/apply 'apply) 2)))))
-	 (maybe-drop))
-
-	((<ghil-mv-call> env src producer consumer)
-	 ;; CONSUMER
-         ;; PRODUCER
-         ;; (mv-call MV)
-         ;; ([tail]-call 1)
-         ;; goto POST
-         ;; MV: [tail-]call/nargs
-         ;; POST: (maybe-drop)
-         (let ((MV (make-label)) (POST (make-label)))
-           (comp-push consumer)
-           (comp-push producer)
-           (push-code! src (make-glil-mv-call 0 MV))
-           (push-code! src (make-glil-call (if tail 'goto/args 'call) 1))
-           (cond ((not tail)
-                  (push-branch! #f 'br POST)))
-           (push-label! MV)
-           (push-code! src (make-glil-call (if tail 'goto/nargs 'call/nargs) 0))
-           (cond ((not tail)
-                  (push-label! POST)
-                  (maybe-drop)))))
-
-        ((<ghil-reified-env> env src)
-         (return-object! src (ghil-env-reify env)))))
 
+(define (codegen x)
+  (define stack '())
+  (define (push-code! src code)
+    (set! stack (cons code stack))
+    (if src (set! stack (cons (make-glil-source src) stack))))
+  (define (var->binding var)
+    (list (ghil-var-name var) (ghil-var-kind var) (ghil-var-index var)))
+  (define (push-bindings! src vars)
+    (if (not (null? vars))
+        (push-code! src (make-glil-bind (map var->binding vars)))))
+  (define (comp tree tail drop)
+    (define (push-label! label)
+      (push-code! #f (make-glil-label label)))
+    (define (push-branch! src inst label)
+      (push-code! src (make-glil-branch inst label)))
+    (define (push-call! src inst args)
+      (for-each comp-push args)
+      (push-code! src (make-glil-call inst (length args))))
+    ;; possible tail position
+    (define (comp-tail tree) (comp tree tail drop))
+    ;; push the result
+    (define (comp-push tree) (comp tree #f #f))
+    ;; drop the result
+    (define (comp-drop tree) (comp tree #f #t))
+    ;; drop the result if unnecessary
+    (define (maybe-drop)
+      (if drop (push-code! #f *ia-drop*)))
+    ;; return here if necessary
+    (define (maybe-return)
+      (if tail (push-code! #f *ia-return*)))
+    ;; return this code if necessary
+    (define (return-code! src code)
+      (if (not drop) (push-code! src code))
+      (maybe-return))
+    ;; return void if necessary
+    (define (return-void!)
+      (return-code! #f *ia-void*))
+    ;; return object if necessary
+    (define (return-object! src obj)
+      (return-code! src (make-glil-const obj)))
     ;;
-    ;; main
-    (record-case ghil
+    ;; dispatch
+    (record-case tree
+      ((<ghil-void>)
+       (return-void!))
+
+      ((<ghil-quote> env src obj)
+       (return-object! src obj))
+
+      ((<ghil-ref> env src var)
+       (return-code! src (make-glil-var 'ref env var)))
+
+      ((<ghil-set> env src var val)
+       (comp-push val)
+       (push-code! src (make-glil-var 'set env var))
+       (return-void!))
+
+      ((<toplevel-define> src name exp)
+       (comp-push exp)
+       (push-code! src (make-glil-var 'define env var))
+       (return-void!))
+
+      ((<conditional> src test then else)
+       ;;     TEST
+       ;;     (br-if-not L1)
+       ;;     THEN
+       ;;     (br L2)
+       ;; L1: ELSE
+       ;; L2:
+       (let ((L1 (make-label)) (L2 (make-label)))
+         (comp-push test)
+         (push-branch! src 'br-if-not L1)
+         (comp-tail then)
+         (if (not tail) (push-branch! #f 'br L2))
+         (push-label! L1)
+         (comp-tail else)
+         (if (not tail) (push-label! L2))))
+
+      ((<sequence> src exps)
+       ;; EXPS...
+       ;; TAIL
+       (if (null? exps)
+           (return-void!)
+           (do ((exps exps (cdr exps)))
+      	 ((null? (cdr exps))
+      	  (comp-tail (car exps)))
+             (comp-drop (car exps)))))
+
+      ((<let> src vars vals body)
+       ;; VALS...
+       ;; (set VARS)...
+       ;; BODY
+       (for-each comp-push vals)
+       (push-bindings! src vars)
+       (for-each (lambda (var) (push-code! #f (make-glil-var 'set env var)))
+      	   (reverse vars))
+       (comp-tail body)
+       (push-code! #f (make-glil-unbind)))
+
+      ((<ghil-mv-bind> env src producer vars rest body)
+       ;; VALS...
+       ;; (set VARS)...
+       ;; BODY
+       (let ((MV (make-label)))
+         (comp-push producer)
+         (push-code! src (make-glil-mv-call 0 MV))
+         (push-code! #f (make-glil-const 1))
+         (push-label! MV)
+         (push-code! #f (make-glil-mv-bind (map var->binding vars) rest))
+         (for-each (lambda (var) (push-code! #f (make-glil-var 'set env var)))
+                   (reverse vars)))
+       (comp-tail body)
+       (push-code! #f (make-glil-unbind)))
+
       ((<ghil-lambda> env src vars rest meta body)
-       (let* ((evars (ghil-env-variables env))
-	      (srcs (pick (lambda (v) (eq? (ghil-var-kind v) 'local)) evars))
-	      (exts (pick (lambda (v) (eq? (ghil-var-kind v) 'external)) evars))
-              (nargs (allocate-indices-linearly! vars))
-              (nlocs (allocate-locals! locs body))
-              (nexts (allocate-indices-linearly! exts)))
-	 ;; meta bindings
-         (push-bindings! #f vars)
-         ;; push on definition source location
-         (if src (set! stack (cons (make-glil-source src) stack)))
-	 ;; copy args to the heap if they're marked as external
-	 (do ((n 0 (1+ n))
-	      (l vars (cdr l)))
-	     ((null? l))
-	   (let ((v (car l)))
-	     (case (ghil-var-kind v)
-               ((external)
-                (push-code! #f (make-glil-argument 'ref n))
-                (push-code! #f (make-glil-external 'set 0 (ghil-var-index v)))))))
-	 ;; compile body
-	 (comp body #t #f)
-	 ;; create GLIL
-         (make-glil-program nargs (if rest 1 0) nlocs nexts meta
-                            (reverse! stack)))))))
+       (return-code! src (codegen tree)))
+
+      ((<ghil-inline> env src inline args)
+       ;; ARGS...
+       ;; (INST NARGS)
+       (let ((tail-table '((call . goto/args)
+                           (apply . goto/apply)
+                           (call/cc . goto/cc))))
+         (cond ((and tail (assq-ref tail-table inline))
+                => (lambda (tail-inst)
+                     (push-call! src tail-inst args)))
+               (else
+                (push-call! src inline args)
+                (maybe-drop)
+                (maybe-return)))))
+
+      ((<ghil-values> env src values)
+       (cond (tail ;; (lambda () (values 1 2))
+              (push-call! src 'return/values values))
+             (drop ;; (lambda () (values 1 2) 3)
+              (for-each comp-drop values))
+             (else ;; (lambda () (list (values 10 12) 1))
+              (push-code! #f (make-glil-const 'values))
+              (push-code! #f (make-glil-call 'link-now 1))
+              (push-code! #f (make-glil-call 'variable-ref 0))
+              (push-call! src 'call values))))
+              
+      ((<ghil-values*> env src values)
+       (cond (tail ;; (lambda () (apply values '(1 2)))
+              (push-call! src 'return/values* values))
+             (drop ;; (lambda () (apply values '(1 2)) 3)
+              (for-each comp-drop values))
+             (else ;; (lambda () (list (apply values '(10 12)) 1))
+              (push-code! #f (make-glil-const 'values))
+              (push-code! #f (make-glil-call 'link-now 1))
+              (push-code! #f (make-glil-call 'variable-ref 0))
+              (push-call! src 'apply values))))
+              
+      ((<ghil-call> env src proc args)
+       ;; PROC
+       ;; ARGS...
+       ;; ([tail-]call NARGS)
+       (comp-push proc)
+       (let ((nargs (length args)))
+         (cond ((< nargs 255)
+                (push-call! src (if tail 'goto/args 'call) args))
+               (else
+                (push-call! src 'mark '())
+                (for-each comp-push args)
+                (push-call! src 'list-mark '())
+                (push-code! src (make-glil-call (if tail 'goto/apply 'apply) 2)))))
+       (maybe-drop))
+
+      ((<ghil-mv-call> env src producer consumer)
+       ;; CONSUMER
+       ;; PRODUCER
+       ;; (mv-call MV)
+       ;; ([tail]-call 1)
+       ;; goto POST
+       ;; MV: [tail-]call/nargs
+       ;; POST: (maybe-drop)
+       (let ((MV (make-label)) (POST (make-label)))
+         (comp-push consumer)
+         (comp-push producer)
+         (push-code! src (make-glil-mv-call 0 MV))
+         (push-code! src (make-glil-call (if tail 'goto/args 'call) 1))
+         (cond ((not tail)
+                (push-branch! #f 'br POST)))
+         (push-label! MV)
+         (push-code! src (make-glil-call (if tail 'goto/nargs 'call/nargs) 0))
+         (cond ((not tail)
+                (push-label! POST)
+                (maybe-drop)))))
+
+      ((<ghil-reified-env> env src)
+       (return-object! src (ghil-env-reify env)))))
+
+  ;;
+  ;; main
+  ;;
+
+  ;; analyze vars: partition into args, locs, exts, and assign indices
+  (record-case x
+    ((<ghil-lambda> env src vars rest meta body)
+     (let* ((evars (ghil-env-variables env))
+            (srcs (pick (lambda (v) (eq? (ghil-var-kind v) 'local)) evars))
+            (exts (pick (lambda (v) (eq? (ghil-var-kind v) 'external)) evars))
+            (nargs (allocate-indices-linearly! vars))
+            (nlocs (allocate-locals! locs body))
+            (nexts (allocate-indices-linearly! exts)))
+       ;; meta bindings
+       (push-bindings! #f vars)
+       ;; push on definition source location
+       (if src (set! stack (cons (make-glil-source src) stack)))
+       ;; copy args to the heap if they're marked as external
+       (do ((n 0 (1+ n))
+            (l vars (cdr l)))
+           ((null? l))
+         (let ((v (car l)))
+           (case (ghil-var-kind v)
+             ((external)
+              (push-code! #f (make-glil-argument 'ref n))
+              (push-code! #f (make-glil-external 'set 0 (ghil-var-index v)))))))
+       ;; compile body
+       (comp body #t #f)
+       ;; create GLIL
+       (make-glil-program nargs (if rest 1 0) nlocs nexts meta
+                          (reverse! stack))))))
 
 (define (allocate-indices-linearly! vars)
   (do ((n 0 (1+ n))
diff --git a/module/language/tree-il/optimize.scm b/module/language/tree-il/optimize.scm
index 52baddb08..14460ebab 100644
--- a/module/language/tree-il/optimize.scm
+++ b/module/language/tree-il/optimize.scm
@@ -22,7 +22,14 @@
 (define-module (language tree-il optimize)
   #:use-module (system base syntax)
   #:use-module (language tree-il)
-  #:export (resolve-primitives!))
+  #:use-module (language tree-il inline)
+  #:export (optimize!))
+
+(define (env-module e)
+  (if e (car e) (current-module)))
+
+(define (optimize! x env opts)
+  (expand-primitives! (resolve-primitives! x (env-module env))))
 
 ;; Possible optimizations:
 ;; * constant folding, propagation