[UPDATE]support closuer in define-method

diff --git a/src/closureTrans.ml b/src/closureTrans.ml
index d92c79a..c2f79e6 100644 (file)
--- a/src/closureTrans.ml
+++ b/src/closureTrans.ml
@@ -65,20 +65,29 @@ let rec closure_fv =
     | _ ->
        Set.empty
 
-let wrap_closure =
+let wrap args body =
+  let fv =
+    Set.elements @@ Set.inter (set_of_list args) (closure_fv body) in
+    if fv = [] then
+      body
+    else
+      Let (List.map (fun x->x,Var x) fv,body)
+
+let expr_trans =
   function
       Lambda (args,body) ->
-       let fv =
-         Set.elements @@ Set.inter (set_of_list args) (closure_fv body) in
-       let body' =
-         if fv = [] then
-           body
-         else
-           Let (List.map (fun x->x,Var x) fv,body) in
-         Lambda (args,body')
+         Lambda (args,wrap args body)
     | e ->
        e
 
+let stmt_trans =
+  function
+      Class (name,super,methods) ->
+       Class (name,super,List.map (fun (name,args,body) ->
+                                     (name,args,wrap args body)) methods)
+    | stmt -> 
+       lift_stmt expr_trans stmt
+
 let trans =
-  lift_program wrap_closure
+  List.map stmt_trans
 

diff --git a/test/test_closuretrans.ml b/test/test_closuretrans.ml
index 6074527..9195863 100644 (file)
--- a/test/test_closuretrans.ml
+++ b/test/test_closuretrans.ml
@@ -3,9 +3,21 @@ open Ast
 open Util
 open ClosureTrans
 
+let ok x y =
+  OUnit.assert_equal
+    ~printer:(fun x-> (Std.dump (List.map Ast.to_string_stmt x) ^ "\n"))
+    x y
+
 test closure_with_args =
-    OUnit.assert_equal
-      ~printer:(fun x-> (Std.dump (List.map Ast.to_string_stmt x)))
-      [Define ("f",Lambda (["x"],(Let (["x",Var "x"],Block [Lambda ([],Block [Var "x"])]))))]
-      (trans @@ compile_string "(define (f x) (lambda () x))")
+  ok
+    [Define ("f",Lambda (["x"],(Let (["x",Var "x"],Block [Lambda ([],Block [Var "x"])]))))]
+    (trans @@ compile_string "(define (f x) (lambda () x))")
+
+test closure_with_class =
+  ok [
+    Class ("Foo",("","Object"),
+          ["init",["self"],
+             Let (["self",Var "self"],
+                  Block [Lambda ([],Block [Var "self"])])])] @@
+    trans @@ compile_string "(define-class Foo (Object) ())  (define-method init ((self Foo)) (lambda () self))"