vm and compiler: removed automatic currying, and several other modifications to the language

and modified the emit module to emit a Pop instruction after every top-level expression.

This change was required because the semantics of the language are pretty clear.
Every expression evaluates to something - meaning that, in the corresponding bytecode,
every expression must have exactly a +1 effect on the data stack. I.e. every expression,
when its corresponding bytecode is evaluated, has the effect of pushing something
to the stack. For values that are not used by another expression, this value
must be immediately popped.

Some optimizations could target this area. For example, for top-level expressions,
it is obvious to the compiler that their values will not be used - hence the compiler
can use optimized versions of some instructions (like StoreLocal and StoreGlobal)
to simply never leave the value on the stack, thus saving an extra Pop instruction
(good for performance and code size).

Same thing applies in function bodies, letrec/let/begin bodies, where expressions
whose values are never used may appear.

It may also make sense to introduce registers to the VM, for the purposes of
parameter passing (such that up to a predetermined number of parameters are
progressively passed through registers instead of pushed to the stack).
This would pair well with eliminating unnecessary currying in the byte code.

bin/comp.ml

@@ -1,16 +1,24 @@
 let ( let* ) = Result.bind;;
 
+
+
 (* Try to interpret some test source code. *)
-let some_source = "(define (+ a b) b)
-                   (+ 1 2)";;
+let some_source = "(define (print-three a b c)
+                    (print a)
+                    (print b)
+                    (print c))
+                   (print-three 'hello 'world 'there)
+                   (print 'fuck)";;
 (* I don't have any built-in functions at all rn, so we just use a dummy function *)
 
 let bruh =
-  let* res = Interpreter.Main.interpret_src some_source in
-  match res with
-  | Int x -> Printf.printf "got %d as result\n" x; Ok ()
-  | _ -> Printf.printf "got something else\n" ; Ok ()
-let _ =
-  match bruh with
-  | Error s -> Printf.printf "%s" s
-  | _ -> ()
+  let* vm = Compiler.Emit.compile_src some_source in
+  print_endline "=== PROGRAM DISASSEMBLY";
+  Vm.Types.print_instrs vm.instrs;
+  print_endline "=== PROGRAM OUTPUT";
+  Vm.interpret vm;
+  Ok ()
+
+let _ = match bruh with
+  | Ok _ -> ()
+  | Error s -> print_endline s

bin/dune

@@ -1,4 +1,4 @@
 (executable
  (name comp)
  (public_name ollisp)
- (libraries str unix compiler interpreter))
+ (libraries str unix compiler vm))

lib/compiler/core_ast.ml

@@ -5,6 +5,7 @@ type literal =
   | Int of int
   | Double of float
   | String of string
+  | Symbol of string
   | Nil
   | Cons of literal * literal
 
@@ -16,8 +17,8 @@ type literal =
 type expression =
   | Literal of literal
   | Var of string
-  | Apply of expression * expression
-  | Lambda of string * expression
+  | Apply of expression * expression list
+  | Lambda of string list * string option * expression
   | If of expression * expression * expression
   | Set of string * expression
   | Begin of expression list
@@ -34,42 +35,30 @@ let rec pair_of_def : Syntactic_ast.def -> string * expression =
 and pair_of_binding (s, e) = (s, of_expr e)
 and pair_of_clause (e1, e2) = (of_expr e1, of_expr e2)
 
-and make_lambda args body =
-  match args with
-    (* TODO: gensym here instead of using _ directly *)
-  | [] -> Lambda ("_", of_body body)
-  | x :: [] -> Lambda (x, of_body body)
-  | x :: xs -> Lambda (x, make_lambda xs body)
-
-and make_apply f args =
-  let rec aux f = function
-    | [] -> Apply (f, Literal Nil)
-    | arg :: [] -> Apply (f, arg)
-    | arg :: args -> aux (Apply (f, arg)) args
-  in aux f args
+and make_lambda (args, rest) body =
+  Lambda (args, rest, body)
 
 (* desugars this...
   (let ((x 5) (y 4)) (f x y))
   ... into this...
-  (((lambda (x) (lambda (y) ((f x) y)))  5)  4)
+  ((lambda (x y) (f x y)) 5 4)
  *)
 and make_let bs body =
   let bs = List.map pair_of_binding bs in
-  let rec aux = function
-    | (s, e) :: rest ->
-       Apply (Lambda (s, aux rest), e)
-    | [] -> of_body body in
-  aux bs
+  let args = List.map (fun (s, _) -> s) bs in
+  let es = List.map (fun (_, e) -> e) bs in
+  Apply (Lambda (args, None, body), es)
 
 (* The Core AST does not feature a letrec node. Instead, we desugar letrecs further
    into a let that binds each symbol to nil, then `set!`s them to their real value
    before running the body.
  *)
 and make_letrec bs exprs =
-  let tmp_bs = List.map (fun (s, _) -> (s, Literal Nil)) bs in
+  let tmp_bs = List.map (fun (_, _) -> Literal Nil) bs in
   let setters = List.fold_right (fun (s, e) acc -> (Set (s, e)) :: acc) bs [] in
+  let args = List.map (fun (s, _) -> s) bs in
   let body = Begin ((List.rev setters) @ exprs) in
-  List.fold_right (fun (s, e) acc -> Apply (Lambda (s, acc), e)) tmp_bs body
+  Apply (Lambda (args, None, body), tmp_bs)
 
 (* We convert a body into a regular letrec form.
    A body is defined as a series of definitions followed by a series
@@ -85,10 +74,8 @@ and of_body : Syntactic_ast.body -> expression = function
      let defs = List.map pair_of_def defs in
      make_letrec defs exprs
 
-(* TODO: currently this ignores the "optional" part of the lambda list,
-   fix this *)
-and of_ll : Syntactic_ast.lambda_list -> string list = function
-  | (sl, _) -> sl
+and of_ll : Syntactic_ast.lambda_list -> string list * string option = function
+  | (sl, rest) -> (sl, rest)
 
 and of_literal : Syntactic_ast.literal -> literal  = function
   | LitInt x -> Int x
@@ -96,12 +83,13 @@ and of_literal : Syntactic_ast.literal -> literal  = function
   | LitString x -> String x
   | LitCons (a, b) -> Cons (of_literal a, of_literal b)
   | LitNil -> Nil
+  | LitSymbol s -> Symbol s
   
 and of_expr : Syntactic_ast.expr -> expression = function
   | Literal l -> Literal (of_literal l)
   | Var x -> Var x
-  | Lambda (ll, b) -> make_lambda (of_ll ll) b
-  | Let (bindings, b) -> make_let bindings b
+  | Lambda ((args, rest), b) -> Lambda (args, rest, of_body b)
+  | Let (bindings, b) -> make_let bindings (of_body b)
   | LetRec (bindings, b) -> make_letrec (List.map pair_of_binding bindings) [(of_body b)]
   | Cond (clauses) -> 
      List.fold_right
@@ -111,7 +99,7 @@ and of_expr : Syntactic_ast.expr -> expression = function
   | If (e1, e2, e3) ->
      If (of_expr e1, of_expr e2, of_expr e3)
   | Set (s, e) -> Set (s, of_expr e)
-  | Apply (f, es) -> make_apply (of_expr f) (List.map of_expr es)
+  | Apply (f, es) -> Apply (of_expr f, List.map of_expr es)
 
 
 and of_syntactic : Syntactic_ast.top_level -> top_level = function

lib/compiler/emit.ml

@@ -7,7 +7,7 @@ type instr = Vm.Types.instr
 
 type pre_instr =
   | Instr of instr
-  | BackPatchMkClosure
+  | BackPatchMkClosure of int
   | BackPatchJumpF
 
 type program = {
@@ -27,8 +27,8 @@ let current_index p =
 let set_instr p i ins =
   Dynarray.set p.instrs i (Instr ins)
 
-let emit_mkclosure p =
-  Ok (Dynarray.add_last p.instrs BackPatchMkClosure)
+let emit_mkclosure p i =
+  Ok (Dynarray.add_last p.instrs  (BackPatchMkClosure i))
 let emit_jumpf p =
   Ok (Dynarray.add_last p.instrs BackPatchJumpF)
 
@@ -47,6 +47,7 @@ let rec compile_one p = function
   | Literal Nil -> emit_constant p (Vm.Types.Nil)
   | Literal (Double x) -> emit_constant p (Vm.Types.Double x)
   | Literal (String s) -> emit_constant p (Vm.Types.String s)
+  | Literal (Symbol s) -> emit_constant p (Vm.Types.Symbol s)
   | Literal (Cons (a, b)) ->
      let* _ = compile_one p (Literal a) in
      let* _ = compile_one p (Literal b) in
@@ -61,12 +62,12 @@ let rec compile_one p = function
   | Set (Global i, expr) ->
      let* _ = compile_one p expr in
      emit_instr p (Vm.Types.StoreGlobal i)
-  | Apply (f, arg) ->
+  | Apply (f, args) ->
      let* _ = compile_one p f in
-     let* _ = compile_one p arg in
-     emit_instr p Vm.Types.Apply
-  | Lambda body ->
-     let* _ = emit_mkclosure p in
+     let* _ = compile_all_no_pop p args in
+     emit_instr p (Vm.Types.Apply (List.length args))
+  | Lambda (arg_count, body) ->
+     let* _ = emit_mkclosure p arg_count in
      Ok (Queue.push ((Dynarray.length p.instrs) - 1, body) p.backpatch)
   | If (test, t, f) ->
 (* *)
@@ -90,10 +91,18 @@ let rec compile_one p = function
      compile_one p e1
   | Begin (e1 :: e2 :: rest) ->
      let* _ = compile_one p e1 in
+     let* _ = emit_instr p Vm.Types.Pop in
      compile_one p (Begin (e2 :: rest))
 
 and compile_all p exprs =
-  Util.traverse (compile_one p) exprs
+  Util.traverse
+    (fun e ->
+      let* _ = compile_one p e in
+      emit_instr p Pop) exprs
+and compile_all_no_pop p exprs =
+  Util.traverse
+    (fun e ->
+      let* _ = compile_one p e in Ok ()) exprs
 
 (* Once we have compiled the top-level expressions, we must now compile
    all of the lambdas we held off on. Some of these will hold more
@@ -101,14 +110,27 @@ and compile_all p exprs =
    of the backpatch queue. 
  *)
 let backpatch_one p (i, b) =
-  Dynarray.set p.instrs i (Instr (MakeClosure (current_index p)));
-  let* _ = compile_one p b in
-  emit_instr p End
-let backpatch p =
+  match Dynarray.get p.instrs i with
+  | BackPatchMkClosure arg_count ->
+     Dynarray.set p.instrs i (Instr (MakeClosure (arg_count, current_index p)));
+     let* _ = compile_one p b in
+     emit_instr p End
+  | _ -> failwith "Can't backpatch anything other than a MakeClosure after compilation"
+let rec backpatch p =
   if Queue.is_empty p.backpatch then
     Ok ()
-  else backpatch_one p (Queue.pop p.backpatch)
+  else
+    (let* _ = backpatch_one p (Queue.pop p.backpatch) in
+     backpatch p)
 
+
+let print_instr = function
+  | Instr i -> Vm.Types.print_one i
+  | BackPatchJumpF ->  "BACKPATCH JUMPF\n"
+  | BackPatchMkClosure i -> "BACKPATCH CLOSURE \n" ^ (string_of_int i)
+let print_instrs =
+  Array.mapi_inplace (fun i ins ->
+    print_endline (Printf.sprintf "%d: %s" i (print_instr ins)); ins)
 let smooth_one = function
   | Instr i -> i
   | _ -> failwith "backpatching process was not complete!"
@@ -123,6 +145,11 @@ let compile (exprs : expression list) (tbl : int SymbolTable.t) =
       backpatch=Queue.create ();
     } in
   let* _ = compile_all program exprs in
+  let* _ = emit_instr program End in
   let* _ = backpatch program in
   let final_instrs = smooth_instrs program in
-  Ok (Vm.Types.make_vm final_instrs (Dynarray.to_array program.constants) (SymbolTable.cardinal tbl))
+  Ok (Vm.make_vm final_instrs (Dynarray.to_array program.constants) ((SymbolTable.cardinal tbl) + 1))
+
+let compile_src src =
+  let* (exprs, tbl) = Scope_analysis.of_src src in
+  compile exprs tbl

lib/compiler/scope_analysis.ml

@@ -33,8 +33,8 @@ type variable =
 type expression =
   | Literal of literal
   | Var of variable
-  | Apply of expression * expression
-  | Lambda of expression
+  | Apply of expression * expression list
+  | Lambda of int * expression
   | If of expression * expression * expression
   | Set of variable * expression
   | Begin of expression list
@@ -66,7 +66,7 @@ let default_global_table =
    be kept at runtime.
  *)
 let extract_globals (top : Core_ast.top_level list) =
-  let id_counter = (ref (-1)) in
+  let id_counter = (ref (SymbolTable.cardinal default_global_table)) in
   let id () =
     id_counter := !id_counter + 1; !id_counter in
   let rec aux tbl = function
@@ -93,11 +93,13 @@ let resolve_global tbl sym =
    global table if we can't find it in the local environment
  *)
 let resolve_symbol tbl env sym =
-  let rec aux counter = function
+  let rec aux counter env_num = function
     | [] -> resolve_global tbl sym
-    | x :: _ when String.equal x sym -> Ok (Local counter)
-    | _ :: rest -> aux (counter + 1) rest
-  in aux 0 env
+    | x :: rest ->
+       match List.find_index (String.equal sym) x with
+       | Some i -> Ok (Local (counter + i))
+       | None  -> aux (counter + (List.length (x :: rest))) (env_num + 1) rest
+  in aux 0 0 env
 
 let resolve_var tbl env sym =
   let* sym = resolve_symbol tbl env sym in
@@ -106,6 +108,16 @@ let resolve_var tbl env sym =
 let resolve_set tbl env sym expr =
   let* sym = resolve_symbol tbl env sym in
   Ok (Set (sym, expr))
+
+let extract_function = function
+  | Core_ast.Define (s, Core_ast.Lambda (args, rest, _)) -> Some (s, args, rest)
+  | _ -> None
+
+let extract_functions exprs =
+  let fs = List.filter Option.is_some (List.map extract_function exprs) in
+  let fs = List.map Option.get fs in
+  List.fold_left (fun t (s, args, rest) -> SymbolTable.add s (args, rest) t) SymbolTable.empty fs
+
 (* We need to do some more sophisticated analysis to detect cases where
    a symbol is accessed before it is defined.
    If a symbol is accessed in a lambda body, that is fine, since that computation
@@ -133,12 +145,15 @@ let convert program =
     | Set (sym, expr) ->
        let* inner = analyze tbl current expr in
        resolve_set tbl current sym inner
-    | Lambda (s, body) ->
-       let* body = (analyze global_tbl (s :: current) body) in
-       Ok (Lambda body)
-    | Apply (f, e) ->
+    | Lambda (args, rest, body) ->
+       let args = (match rest with
+                  | Some s -> List.append args [s]
+                  | None -> args) in
+       let* body = (analyze global_tbl (args :: current) body) in
+       Ok (Lambda (List.length args, body))
+    | Apply (f, es) ->
        let* f = analyze tbl current f in
-       let* e = analyze tbl current e in
+       let* e = Util.traverse (analyze tbl current) es in
        Ok (Apply (f, e))
     | If (test, pos, neg) ->
        let* test = analyze tbl current test in
@@ -156,7 +171,7 @@ let convert program =
        let tbl = SymbolTable.add s (SymbolTable.find s global_tbl) tbl in
        (analyze tbl [] (Set (s, e))) :: (aux tbl rest)
   in
-  let* program = traverse (fun x -> x) (aux SymbolTable.empty program) in
+  let* program = traverse (fun x -> x) (aux default_global_table program) in
   Ok (program, global_tbl)
 
 let of_src src =

lib/compiler/syntactic_ast.ml

@@ -9,6 +9,7 @@ type literal =
   | LitDouble of float
   | LitString of string
   | LitCons of literal * literal
+  | LitSymbol of string
   | LitNil
 
 type lambda_list = string list * string option
@@ -81,13 +82,14 @@ let expect_sym = function
    TODO: add diagnostics, e.g. what sexpr, specifically, couldn't be turned to a list?
    generally more debugging is needed in this module.
  *)
-let rec list_of_sexpr = function
+let rec  list_of_sexpr = function
   | LCons (i, next) ->
      let* next = list_of_sexpr next in
      Ok (i :: next)
   | LNil -> Ok []
   | _ -> Error "cannot transform sexpr into list, malformed sexpr!"
 
+  
 (* parse the argument list of a lambda form *)
 let parse_lambda_list cons =
   let rec aux acc = function
@@ -203,6 +205,18 @@ and builtin_set cons =
   let* expr = unwrap_exp (transform expr) in
   exp (Set (sym, expr))
 
+and builtin_quote cons =
+  let* expr = sexpr_cadr cons in
+  let lit x = exp (Literal x) in
+  let rec aux = function
+    | LSymbol s -> (LitSymbol s)
+    | LInt x -> (LitInt x)
+    | LDouble x -> (LitDouble x)
+    | LString x -> (LitString x)
+    | LCons (a, b) -> (LitCons (aux a, aux b))
+    | LNil -> (LitNil) in
+  lit (aux expr)
+
 and apply f args =
   let* args = list_of_sexpr args in
   let* args = traverse (fun x -> unwrap_exp (transform x)) args in
@@ -217,6 +231,7 @@ and builtin_symbol = function
   | "cond" -> builtin_cond
   | "if" -> builtin_if
   | "set!" -> builtin_set
+  | "quote" -> builtin_quote
   | _ -> (function
        | LCons (f, args) -> apply f args
        | _ -> Error "Invalid function application!")
@@ -236,7 +251,8 @@ and transform : lisp_ast -> (top_level, string) result = function
 let make (expr : Parser.Ast.lisp_ast) : (top_level, string) result =
   transform expr
 
-
+let of_src s =
+  Util.traverse make (Parser.parse_str s)
 
 (* Printing, for debug purposes *)
 let pf = Printf.sprintf
@@ -266,7 +282,7 @@ and print_literal = function
   | LitString x -> pf "\"%s\"" x
   | LitNil -> pf "nil"
   | LitCons (a, b) -> pf "(%s . %s)" (print_literal a) (print_literal b)
-
+  | LitSymbol s -> pf "'%s" s
 and print_expr = function
   | Literal l -> print_literal l
   | Lambda (ll, (defs, exprs)) ->

lib/interpreter/dune

@@ -1,3 +0,0 @@
-(library
- (name interpreter)
- (libraries compiler))

lib/interpreter/main.ml

@@ -1,76 +0,0 @@
-
-let ( let* ) = Result.bind
-let traverse = Compiler.Util.traverse
-
-type runtime_value =
-  | Int of int
-  | Double of float
-  | String of string
-  | Nil
-  | Cons of runtime_value * runtime_value
-  | Symbol of string
-(* The rest can't appear as literal values, and are constructed in other ways *)
-  | Closure of Compiler.Scope_analysis.expression * (runtime_value ref list)
-
-let rec interpret_literal = function
-  | Compiler.Core_ast.Int x -> Ok (Int x)
-  | Double x -> Ok (Double x)
-  | String s -> Ok (String s)
-  | Cons (a, b) ->
-     let* a = interpret_literal a in
-     let* b = interpret_literal b in
-     Ok (Cons (a, b))
-  | Nil -> Ok (Nil)
-    
-
-let rec interpret_one expr env globals =
-  match expr with
-  | Compiler.Scope_analysis.Literal l -> interpret_literal l
-  | Var (Local i) ->
-     (match (List.nth_opt env i) with
-      | None -> Error "Error while accessing local variable!"
-      | Some x -> Ok !x)
-  | Var (Global i) ->
-     Ok (Array.get globals i)
-  | Apply (f, e) ->
-     let* f = interpret_one f env globals in
-     let* e = interpret_one e env globals in
-     (match f with
-      | Closure (body, inner_env) ->
-         let f_env = (ref e) :: inner_env in
-         interpret_one body f_env globals
-      | _ -> Error "Cannot apply an argument to non-closure value!")
-  | Lambda body ->
-     Ok (Closure (body, env))
-  | If (test, then_e, else_e) ->
-     let* test = interpret_one test env globals in
-     (match test with
-      | Nil -> interpret_one else_e env globals
-      | _ -> interpret_one then_e env globals)
-  | Set ((Local i), e) ->
-     (match (List.nth_opt env i) with
-      | None -> Error "Error while setting local variable!"
-      | Some r ->
-         let* e = interpret_one e env globals in
-         r := e; Ok e)
-  | Set ((Global i), e) ->
-     let* e = interpret_one e env globals in
-     Array.set globals i e; Ok e
-  | Begin [] -> Ok Nil
-  | Begin [e] -> interpret_one e env globals
-  | Begin (e :: rest) ->
-     let* e = interpret_one e env globals in
-     ignore e; interpret_one (Begin rest) env globals
-
-let interpret program global_syms =
-  let count = Compiler.Scope_analysis.SymbolTable.cardinal global_syms in
-  let globals : runtime_value array = Array.make count Nil in
-  
-  interpret_one (Begin program) [] globals
-
-
-
-let interpret_src src =
-  let* (program, globals) = Compiler.Scope_analysis.of_src src in
-  interpret program globals
-

lib/vm/native.ml

@@ -3,12 +3,11 @@
    Stuff like printing to the screen, file I/O etc will be implemented
    here.
  *)
-
 open Types
 
-let builtin_print (v : value) =
-  ignore v;
-  Nil
+let builtin_print (v : Types.value ref list) =
+  List.iter (fun r -> print_endline (print_value !r)) v;
+  Types.Nil
 
 let table = [|
     builtin_print
@@ -16,4 +15,3 @@ let table = [|
 
 
   
-  

lib/vm/types.ml

@@ -6,7 +6,7 @@ type value =
   | Nil
   | Cons of value * value
   | Symbol of string
-  | Closure of int * value ref list
+  | Closure of int * int * value ref list
   | Native of int (* This is basically a syscall, each ID represents a primitive operation
                      that should have a well-defined effect. These will be further detailed
                      in the language documentation
@@ -20,8 +20,8 @@ type instr =
   | StoreGlobal of int
   | MakeCons
   | Pop (* discards top of stack *)
-  | Apply
-  | MakeClosure of int
+  | Apply of int (* arg count *)
+  | MakeClosure of int * int (* arg count, code pointer *)
   | Jump of int
   | JumpF of int (* jump if false. *)
   | End
@@ -33,18 +33,42 @@ type vm_state = {
     globals : value array;
     constants : value array;
     mutable env : value ref list;
-    mutable stack : value list
+    mutable stack : value list;
+    mutable call_stack : (int * (value ref list)) list;
   }
 
 
-let make_vm instrs constants global_count =
-  {
-    i = 0;
-    instrs = instrs;
-    globals = Array.make global_count Nil;
-    constants = constants;
-    env = [];
-    stack = [];
-  }
+let p = Printf.sprintf 
 
-(* TODO: add facilities to print the VM state in case of errors. *)
+let rec print_value = function
+    | Int x -> p "%d" x
+    | Double x -> p "%f" x
+    | String x -> p "\"%s\"" x
+    | Nil -> p "'()"
+    | Cons (a, b) -> p "(%s . %s)" (print_value a) (print_value b)
+    | Symbol x -> p "'%s" x
+    | Closure (a, i, _) -> p "<closure of %d args at %d>" a i
+    | Native i -> p "<native %d>" i
+
+
+let print_one = function
+    | Constant i -> p "CONSTANT %d\n" i
+    | LoadLocal i -> p "LOCAL %d\n" i
+    | LoadGlobal i -> p "GLOBAL %d\n" i
+    | StoreLocal i -> p "STORE_LOCAL %d\n" i
+    | StoreGlobal i -> p "STORE_GLOBAL %d\n" i
+    | MakeCons -> p "CONS\n"
+    | Pop -> p "POP\n"
+    | Apply i -> p "APPLY %d\n" i
+    | MakeClosure (a, i) -> p "MKCLOSURE %d, %d\n" a i
+    | Jump i -> p "JMP %d\n" i
+    | JumpF i -> p "JMPF %d\n" i
+    | End -> p "END\n"
+    | NOOP -> p "NOOP\n"
+
+let print_instrs instrs =  
+  Array.mapi_inplace
+    (fun i ins ->
+      print_string (p "%d: %s" i (print_one ins));
+      ins)
+    instrs

lib/vm/vm.ml

@@ -15,48 +15,83 @@ let set_local state i v =
 let pop_one state =
   match state.stack with
   | v :: rest -> state.stack <- rest; v
-  | [] -> failwith "VM error: cannot pop from empty stack!"
+  | [] -> failwith ("VM error: cannot pop from empty stack! " )
+let pop_args state count =
+  let rec aux acc i =
+    if i <= 0 then acc
+    else aux ((ref (pop_one state)) :: acc) (i - 1)
+  in aux [] count
+let peek_one state =
+  match state.stack with
+  | v :: _ -> v
+  | [] -> failwith ("VM error: cannot peek on empty stack! " )
 
 let push state v =
-  state.stack <- v :: state.stack
+  state.stack <- (v :: state.stack)
 
-let rec do_apply state =
+let trace state =
+  let stack () = List.fold_left (fun acc x -> acc ^ " " ^ (Types.print_value x)) "" state.stack in
+  let env () = List.fold_left (fun acc x -> acc ^ " " ^ (Types.print_value !x)) "" state.env in
+  Printf.printf "%d: \n\tstack: [%s ]\n\tenv:[%s]\n" state.i (stack ()) (env ())
+
+let rec do_apply state arg_count =
   let cur_env = state.env in
   let cur_i = state.i in
-  let arg = pop_one state in
+  let args = pop_args state arg_count in
   let f = pop_one state in
   match f with
-  | Closure (x, e) ->
-     state.env <- ref arg :: e;
+  | Closure (a, _, _) when a != arg_count -> failwith "Wrong argument count to function"
+  | Closure (_, x, e) ->
+     state.call_stack <- (cur_i, cur_env) :: state.call_stack;
      state.i <- x;
-     interpret state;
-     state.env <- cur_env;
-     state.i <- cur_i
+     state.env <- List.append args e;
+     interpret state
   | Native x ->
-     push state (Native.table.(x) arg)
+     push state (Native.table.(x) args);
+     interpret state
   | _ -> failwith "Cannot apply non-closure object"
 
 and interpret state =
+  trace state;
   let i = state.i in
   state.i <- i + 1;
-  (match state.instrs.(state.i) with
+  (match state.instrs.(i) with
   | Constant x -> push state state.constants.(x) ; interpret state
   | LoadLocal x -> push state (load_local state x) ; interpret state
   | LoadGlobal x -> push state state.globals.(x) ; interpret state
-  | StoreLocal x -> set_local state x (pop_one state) ; interpret state
-  | StoreGlobal x -> Array.set state.globals x (pop_one state) ; interpret state
+  | StoreLocal x -> set_local state x (peek_one state) ; interpret state
+  | StoreGlobal x -> Array.set state.globals x (peek_one state) ; interpret state
   | MakeCons ->
      let cdr = pop_one state in
      let car = pop_one state in
      push state (Cons (car, cdr))
   | Pop -> ignore (pop_one state) ; interpret state
-  | Apply -> do_apply state ; interpret state
-  | MakeClosure x -> push state (Closure (x, state.env)); interpret state
+  | Apply a -> do_apply state a
+  | MakeClosure (args, x) -> push state (Closure (args, x, state.env)); interpret state
   | Jump target -> state.i <- target ; interpret state
   | JumpF target ->
      (match (pop_one state) with
      | Nil -> state.i <- target
      | _ -> ()); interpret state 
-  | End -> ()
+  | End ->
+     (match state.call_stack with
+     | [] ->
+        print_endline "\nPROGRAM HAS SUCCESSFULLY TERMINATED"
+     | (old_i, old_env) :: rest ->
+        state.call_stack <- rest;
+        state.env <- old_env;
+        state.i <- old_i;
+        interpret state)
   | NOOP -> interpret state)
 
+let make_vm instrs constants global_count =
+  let globals = Array.init global_count (fun x -> if x < (Array.length Native.table) then Native x else Nil) in
+  {
+    i = 0;
+    instrs = instrs;
+    globals = globals;
+    constants = constants;
+    env = [];
+    stack = [];
+    call_stack = [];
+  }