Changed constant node types, added VarDecs|LocalFuns wrappers, implemented first steps of assembler

Makefile

@@ -1,6 +1,6 @@
 RESULT := civicc
 PHASES := load parse print desug context_analysis expand_dims typecheck \
-	dim_reduce bool_op extern_vars constant_propagation
+	dim_reduce bool_op extern_vars constant_propagation depth_analysis assemble
 SOURCES := types.ml stringify.mli stringify.ml util.mli util.ml lexer.mll \
 	parser.mly $(patsubst %,phases/%.ml,$(PHASES)) main.ml
 PRE_TARGETS := types.cmi types.o stringify.cmi stringify.o util.cmi util.o

README.md

@@ -9,3 +9,5 @@ Issues & TODO
 
 - Keep file content in buffer to prevent error messages from crashing when
   reading from stdin.
+- Maybe declarations should be saved as references to allow the declarations to
+  be modified (does make pattern matching a bit harder).

main.ml

@@ -12,28 +12,30 @@ let compile () =
     in
     run_phases Empty [
         Load.phase;
-        Print.phase;
+        (*Print.phase;*)
         Parse.phase;
-        Print.phase;
+        (*Print.phase;*)
         Desug.phase;
-        Print.phase;
+        (*Print.phase;*)
         Context_analysis.phase;
-        Print.phase;
+        (*Print.phase;*)
         Typecheck.phase;
-        Print.phase;
+        (*Print.phase;*)
         Expand_dims.phase;
-        Print.phase;
+        (*Print.phase;*)
         Bool_op.phase;
-        Print.phase;
+        (*Print.phase;*)
         Dim_reduce.phase;
-        Print.phase;
+        (*Print.phase;*)
         Extern_vars.phase;
-        Print.phase;
+        (*Print.phase;*)
         Constant_propagation.phase;
         Print.phase;
-        (*
+        Depth_analysis.phase;
+        Print.phase;
         Assemble.phase;
         Print.phase;
+        (*
         Peephole.phase;
         Print.phase;
         *)

parser.mly

@@ -117,7 +117,7 @@ param:
 
 fun_body:
     | var_dec* local_fun_dec* statement* loption(return_statement)
-    { $1 @ $2 @ $3 @ $4 }
+    { VarDecs $1 :: (LocalFuns $2) :: $3 @ $4 }
 
 return_statement:
     (* return statement: use location of return value *)
@@ -177,7 +177,7 @@ statement:
     | FOR; LPAREN; INT; cnt=ID; ASSIGN; start=expr; COMMA; stop=expr; RPAREN;
       body=block
     { let loc = loc $startpos(cnt) $endpos(cnt) in
-        For (cnt, start, stop, IntConst (1, []), Block body, loc) }
+        For (cnt, start, stop, Const (IntVal 1, []), Block body, loc) }
 
     | FOR; LPAREN; INT; cnt=ID; ASSIGN; start=expr; COMMA; stop=expr; COMMA;
       step=expr; RPAREN; body=block
@@ -198,9 +198,9 @@ expr:
     | SUB; expr                         { Monop (Neg, $2, loc $startpos $endpos) }    %prec NEG
     | NOT; expr                         { Monop (Not, $2, loc $startpos $endpos) }
     | LPAREN; basic_type; RPAREN; expr  { TypeCast ($2, $4, loc $startpos $endpos) }  %prec CAST
-    | FLOAT_CONST                       { FloatConst ($1, loc $startpos $endpos) }
-    | INT_CONST                         { IntConst ($1, loc $startpos $endpos) }
-    | BOOL_CONST                        { BoolConst ($1, loc $startpos $endpos) }
+    | FLOAT_CONST                       { Const (FloatVal $1, loc $startpos $endpos) }
+    | INT_CONST                         { Const (IntVal $1, loc $startpos $endpos) }
+    | BOOL_CONST                        { Const (BoolVal $1, loc $startpos $endpos) }
     | ID; array_const                   { Var ($1, Some $2, loc $startpos $endpos) }
     | array_const                       { ArrayConst ($1, loc $startpos $endpos) }
 

phases/assemble.ml

@@ -1,9 +1,6 @@
 open Types
 open Util
 
-let store ctype = function
-    | 0 -> StoreGlob ()
-
 let assemble program =
     let labcounter = ref 0 in
     let genlabel suffix =
@@ -11,63 +8,77 @@ let assemble program =
         string_of_int !labcounter ^ "_" ^ suffix
     in
 
-    let consts = ref [] in
-    let const_index const =
-        let rec trav_consts i = function
-            | [] -> consts := !consts @ [const]; i
-            | hd :: _ when hd = const -> i
-            | hd :: tl -> trav_consts (i + 1) tl
-        in
-        trav_consts 0 !consts
-    in
+    let consts = Hashtbl.create 20 in
 
-    let rec trav node =
+    let rec trav_args callstack localfuns node =
+        let trav = trav_args callstack localfuns in
         let rec trav_all = function
             | [] -> []
             | hd :: tl -> trav hd @ (trav_all tl)
         in
+        let rec traverse_localfuns = function
+            | LocalFuns funs -> trav_all funs
+            | Block body -> List.concat (List.map traverse_localfuns body)
+            | _ -> []
+        in
         match node with
         | Program (decls, _) ->
             trav_all decls
 
         | FunDec (ret_type, name, params, _) ->
-            [Import (name, ret_type, List.map ctypeof params)]
+            [Import (name, ret_type, List.map typeof params)]
 
         | FunDef (export, ret_type, name, params, body, _) ->
-            let label = name in
-            let param_types = List.map ctypeof params in
-            let export = match export with
-                | false -> []
-                | true -> [Export (name, ret_type, param_types, label)]
-            in
-            Comment ("function \"" ^ name ^ "\":") ::
-            (export @ (Label label :: (trav body)))
+            localfuns := node :: !localfuns;
+            let callstack = name :: callstack in
+            let localfuns = ref [] in
+            let label = String.concat "$" (List.rev callstack) in
+            (if export then
+                let param_types = List.map typeof params in
+                [Export (name, ret_type, param_types, label)]
+            else []) @
+            [
+                Comment ("function \"" ^ label ^ "\":");
+                Label label;
+                RtnEnter (indexof node);
+            ] @
+            (trav_args callstack localfuns body) @
+            (match ret_type with Void -> [Ret Void] | _ -> []) @
+            [EmptyLine] @
+            (traverse_localfuns body)
 
-        | VarDec (ctype, name, None, _) ->
-            []
+        (* Local fucntions are traversed elsewhere *)
+        | LocalFuns _ -> []
 
-        | VarLet (Assign (name, None, value, _), ctype, depth) ->
-            [store ctype depth]
-        (*
-        | VarLet (Assign (name, Some indices, value, _), ctype, depth) ->
-            [store deoth]
-        *)
+        | Block body -> trav_all body
 
-        | BoolConst _ ->
-            [LoadImm node]
+        | VarLet (dec, None, value, _) ->
+            let store = match (depthof dec, depthof node) with
+                | (0, _)            -> StoreGlob (typeof dec, indexof dec)
+                | (a, b) when a = b -> StoreLoc  (typeof dec, indexof dec)
+                | (a, b)            -> StoreRel  (typeof dec, b - a, indexof dec)
+            in
+            trav value @ [store]
 
-        | IntConst _ | FloatConst _ ->
-            [LoadConst (ctypeof node, const_index node)]
+        | Return (value, _) ->
+            trav value @ [Ret (typeof value)]
+
+        | Const (BoolVal _, _) ->
+            [LoadImm node]
+        | Const (value, _) ->
+            let def = if Hashtbl.mem consts value then [] else (
+                Hashtbl.add consts value true;
+                [ConstDef value]
+            ) in
+            def @ [LoadConst (typeof node, indexof node)]
 
         | _ -> []
         (*| _ -> raise InvalidNode*)
     in
-    let instrs = trav program in
-    let const_defs = List.map (fun c -> Const c) !consts in
-    const_defs @ instrs
+    trav_args [] (ref []) program
 
 let rec phase input =
-    prerr_endline "- Assembly";
+    log_line 2 "- Assembly";
     match input with
-    | Types node -> Assembly (assemble node)
+    | Ast node -> Assembly (assemble node)
     | _ -> raise (InvalidInput "assembly")

phases/bool_op.ml

@@ -21,9 +21,9 @@ open Util
 
 let cast ctype node = TypeCast (ctype, node, [Type ctype])
 
-let boolconst  value = BoolConst  (value, [Type Bool])
-let intconst   value = IntConst   (value, [Type Int])
-let floatconst value = FloatConst (value, [Type Float])
+let boolconst  value = Const (BoolVal  value, [Type Bool])
+let intconst   value = Const (IntVal   value, [Type Int])
+let floatconst value = Const (FloatVal value, [Type Float])
 
 let rec trav_binop = function
     | ((Eq | Ne) as op, left, right, loc) ->
@@ -62,5 +62,5 @@ and bool_op = function
 let rec phase input =
     log_line 2 "- Convert bool operations";
     match input with
-    | Types node -> Types (bool_op node)
+    | Ast node -> Ast (bool_op node)
     | _ -> raise (InvalidInput "bool operations")

phases/constant_propagation.ml

@@ -16,77 +16,75 @@ open Util
 let is_const_name name =
     Str.string_match (Str.regexp "^.+\\$\\$[0-9]+$") name 0
 
-let is_const = function
-    | BoolConst _ | IntConst _ | FloatConst _ -> true
-    | _ -> false
+let is_const = function Const _ -> true | _ -> false
 
 let eval_monop = function
-    | (Not, BoolConst  (value, _), ann) -> BoolConst  (not value, ann)
-    | (Neg, IntConst   (value, _), ann) -> IntConst   (-value, ann)
-    | (Neg, FloatConst (value, _), ann) -> FloatConst (-.value, ann)
+    | (Not, Const (BoolVal  value, _), ann) -> Const (BoolVal  (not value), ann)
+    | (Neg, Const (IntVal   value, _), ann) -> Const (IntVal   (-value), ann)
+    | (Neg, Const (FloatVal value, _), ann) -> Const (FloatVal (-.value), ann)
     | (op, opnd, ann) -> Monop (op, opnd, ann)
 
 let eval_binop = function
     (* Arithmetic *)
-    | (Add, IntConst (left, _), IntConst (right, _), ann) ->
-        IntConst (left + right, ann)
-    | (Add, FloatConst (left, _), FloatConst (right, _), ann) ->
-        FloatConst (left +. right, ann)
+    | (Add, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (IntVal (left + right), ann)
+    | (Add, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (FloatVal (left +. right), ann)
 
-    | (Sub, IntConst (left, _), IntConst (right, _), ann) ->
-        IntConst (left - right, ann)
-    | (Sub, FloatConst (left, _), FloatConst (right, _), ann) ->
-        FloatConst (left -. right, ann)
+    | (Sub, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (IntVal (left - right), ann)
+    | (Sub, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (FloatVal (left -. right), ann)
 
-    | (Mul, IntConst (left, _), IntConst (right, _), ann) ->
-        IntConst (left * right, ann)
-    | (Mul, FloatConst (left, _), FloatConst (right, _), ann) ->
-        FloatConst (left *. right, ann)
+    | (Mul, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (IntVal (left * right), ann)
+    | (Mul, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (FloatVal (left *. right), ann)
 
-    | (Div, IntConst (left, _), IntConst (right, _), ann) ->
-        IntConst (left / right, ann)
-    | (Div, FloatConst (left, _), FloatConst (right, _), ann) ->
-        FloatConst (left /. right, ann)
+    | (Div, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (IntVal (left / right), ann)
+    | (Div, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (FloatVal (left /. right), ann)
 
-    | (Mod, IntConst (left, _), IntConst (right, _), ann) ->
-        IntConst (left mod right, ann)
+    | (Mod, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (IntVal (left mod right), ann)
 
     (* Relational *)
-    | (Eq, IntConst (left, _), IntConst (right, _), ann) ->
-        BoolConst (left = right, ann)
-    | (Eq, FloatConst (left, _), FloatConst (right, _), ann) ->
-        BoolConst (left = right, ann)
-
-    | (Ne, IntConst (left, _), IntConst (right, _), ann) ->
-        BoolConst (left != right, ann)
-    | (Ne, FloatConst (left, _), FloatConst (right, _), ann) ->
-        BoolConst (left != right, ann)
-
-    | (Gt, IntConst (left, _), IntConst (right, _), ann) ->
-        BoolConst (left > right, ann)
-    | (Gt, FloatConst (left, _), FloatConst (right, _), ann) ->
-        BoolConst (left > right, ann)
-
-    | (Lt, IntConst (left, _), IntConst (right, _), ann) ->
-        BoolConst (left < right, ann)
-    | (Lt, FloatConst (left, _), FloatConst (right, _), ann) ->
-        BoolConst (left < right, ann)
-
-    | (Ge, IntConst (left, _), IntConst (right, _), ann) ->
-        BoolConst (left >= right, ann)
-    | (Ge, FloatConst (left, _), FloatConst (right, _), ann) ->
-        BoolConst (left >= right, ann)
-
-    | (Le, IntConst (left, _), IntConst (right, _), ann) ->
-        BoolConst (left <= right, ann)
-    | (Le, FloatConst (left, _), FloatConst (right, _), ann) ->
-        BoolConst (left <= right, ann)
+    | (Eq, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (BoolVal (left = right), ann)
+    | (Eq, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (BoolVal (left = right), ann)
+
+    | (Ne, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (BoolVal (left != right), ann)
+    | (Ne, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (BoolVal (left != right), ann)
+
+    | (Gt, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (BoolVal (left > right), ann)
+    | (Gt, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (BoolVal (left > right), ann)
+
+    | (Lt, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (BoolVal (left < right), ann)
+    | (Lt, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (BoolVal (left < right), ann)
+
+    | (Ge, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (BoolVal (left >= right), ann)
+    | (Ge, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (BoolVal (left >= right), ann)
+
+    | (Le, Const (IntVal left, _), Const (IntVal right, _), ann) ->
+        Const (BoolVal (left <= right), ann)
+    | (Le, Const (FloatVal left, _), Const (FloatVal right, _), ann) ->
+        Const (BoolVal (left <= right), ann)
 
     (* Logical *)
-    | (And, BoolConst (left, _), BoolConst (right, _), ann) ->
-        BoolConst (left && right, ann)
-    | (Or, BoolConst (left, _), BoolConst (right, _), ann) ->
-        BoolConst (left || right, ann)
+    | (And, Const (BoolVal left, _), Const (BoolVal right, _), ann) ->
+        Const (BoolVal (left && right), ann)
+    | (Or, Const (BoolVal left, _), Const (BoolVal right, _), ann) ->
+        Const (BoolVal (left || right), ann)
 
     | (op, left, right, ann) -> Binop (op, left, right, ann)
 
@@ -139,22 +137,22 @@ let rec propagate consts node =
         let texp = propagate texp in
         let fexp = propagate fexp in
         (match cond with
-        | BoolConst (value, _) -> if value then texp else fexp
+        | Const (BoolVal value, _) -> if value then texp else fexp
         | _ -> Cond (cond, texp, fexp, ann)
         )
 
     | TypeCast (ctype, value, ann) ->
         let value = propagate value in
         (match (ctype, value) with
-        | (Bool,  BoolConst  (value, _)) -> BoolConst (value, ann)
-        | (Bool,  IntConst   (value, _)) -> BoolConst (value != 1, ann)
-        | (Bool,  FloatConst (value, _)) -> BoolConst (value != 1.0, ann)
-        | (Int,   BoolConst  (value, _)) -> IntConst ((if value then 1 else 0), ann)
-        | (Int,   IntConst   (value, _)) -> IntConst (value, ann)
-        | (Int,   FloatConst (value, _)) -> IntConst (int_of_float value, ann)
-        | (Float, BoolConst  (value, _)) -> FloatConst ((if value then 1. else 0.), ann)
-        | (Float, IntConst   (value, _)) -> FloatConst (float_of_int value, ann)
-        | (Float, FloatConst (value, _)) -> FloatConst (value, ann)
+        | (Bool,  Const (BoolVal  value, _)) -> Const (BoolVal value, ann)
+        | (Bool,  Const (IntVal   value, _)) -> Const (BoolVal (value != 1), ann)
+        | (Bool,  Const (FloatVal value, _)) -> Const (BoolVal (value != 1.0), ann)
+        | (Int,   Const (BoolVal  value, _)) -> Const (IntVal (if value then 1 else 0), ann)
+        | (Int,   Const (IntVal   value, _)) -> Const (IntVal value, ann)
+        | (Int,   Const (FloatVal value, _)) -> Const (IntVal (int_of_float value), ann)
+        | (Float, Const (BoolVal  value, _)) -> Const (FloatVal (if value then 1. else 0.), ann)
+        | (Float, Const (IntVal   value, _)) -> Const (FloatVal (float_of_int value), ann)
+        | (Float, Const (FloatVal value, _)) -> Const (FloatVal value, ann)
         | _ -> TypeCast (ctype, value, ann)
         )
 
@@ -167,8 +165,8 @@ let rec prune_vardecs consts = function
 let rec phase input =
     log_line 2 "- Constant propagation";
     match input with
-    | Types node ->
+    | Ast node ->
         let consts = Hashtbl.create 32 in
         let node = propagate consts node in
-        Types (prune_vardecs consts node)
+        Ast (prune_vardecs consts node)
     | _ -> raise (InvalidInput "constant propagation")

phases/context_analysis.ml

@@ -25,28 +25,25 @@ let check_in_scope name errnode scope =
         in
         raise (NodeError (errnode, msg))
 
-let rec analyse scope depth args node =
-    (* add_to_scope uses args, so it needs to be defined here *)
-    let add_to_scope name dec depth scope =
-        let (vars, funs) = scope in
-        let (name, tbl, name_type) = match name with
-            | Varname name  -> (name, vars, "variable")
-            | Funcname name -> (name, funs, "function")
-        in
-        match mapfind name tbl with
-        (* Identifiers of lower depth may be overwritten, but idenetifiers at
-         * the same depth must be unique for consistency *)
-        | Some (orig, orig_depth, _) when orig_depth >= depth ->
-            let msg = sprintf "Error: cannot redeclare %s \"%s\"" name_type name in
-            prerr_loc_msg (locof dec) msg args.verbose;
-            prerr_loc_msg (locof orig) "Previously declared here:" args.verbose;
-            raise EmptyError
-        | Some _ ->
-            Hashtbl.replace tbl name (dec, depth, name_type)
-        | None ->
-            Hashtbl.add tbl name (dec, depth, name_type)
+let add_to_scope name dec depth (vars, funs) =
+    let (name, tbl, name_type) = match name with
+        | Varname name  -> (name, vars, "variable")
+        | Funcname name -> (name, funs, "function")
     in
+    match mapfind name tbl with
+    (* Identifiers of lower depth may be overwritten, but idenetifiers at
+        * the same depth must be unique for consistency *)
+    | Some (orig, orig_depth, _) when orig_depth >= depth ->
+        let msg = sprintf "Error: cannot redeclare %s \"%s\"" name_type name in
+        prerr_loc_msg (locof dec) msg args.verbose;
+        prerr_loc_msg (locof orig) "Previously declared here:" args.verbose;
+        raise EmptyError
+    | Some _ ->
+        Hashtbl.replace tbl name (dec, depth, name_type)
+    | None ->
+        Hashtbl.add tbl name (dec, depth, name_type)
 
+let rec analyse scope depth node =
     let rec collect node = match node with
         (* Add node reference for this varname to vars map *)
         | VarDec (ctype, name, init, ann) ->
@@ -59,9 +56,12 @@ let rec analyse scope depth args node =
             node
 
         (* For global vars, only the name and array dimensions *)
-        | GlobalDec (Array (ctype, dims), name, ann) ->
-            let node = GlobalDec (Array (ctype, List.map collect dims), name,
-                                  Depth depth :: ann) in
+        | GlobalDec (ctype, name, ann) ->
+            let ctype = match ctype with
+                | Array (ctype, dims) -> Array (ctype, List.map collect dims)
+                | _ -> ctype
+            in
+            let node = GlobalDec (ctype, name, Depth depth :: ann) in
             add_to_scope (Varname name) node depth scope;
             node
 
@@ -70,12 +70,11 @@ let rec analyse scope depth args node =
             add_to_scope (Varname name) node depth scope;
             node
 
-        | GlobalDec (ctype, name, ann) ->
-            let node = GlobalDec (ctype, name, Depth depth :: ann) in
-            add_to_scope (Varname name) node depth scope;
-            node
-
         | GlobalDef (export, ctype, name, init, ann) ->
+            let ctype = match ctype with
+                | Array (ctype, dims) -> Array (ctype, List.map collect dims)
+                | _ -> ctype
+            in
             let node = GlobalDef (export, ctype, name, init, Depth depth :: ann) in
             add_to_scope (Varname name) node depth scope;
             node
@@ -117,7 +116,7 @@ let rec analyse scope depth args node =
             let (vars, funs) = scope in
             let local_scope = (Hashtbl.copy vars, Hashtbl.copy funs) in
             let params = List.map (traverse local_scope (depth + 1)) params in
-            let body = analyse local_scope (depth + 1) args body in
+            let body = analyse local_scope (depth + 1) body in
             FunDef (export, ret_type, name, params, body, ann)
 
         | Param (Array (ctype, dims), name, ann) as node ->
@@ -154,12 +153,12 @@ let rec analyse scope depth args node =
     let node = traverse scope depth node in
     node
 
-let analyse_context args program =
+let analyse_context program =
     let scope = (Hashtbl.create 20, Hashtbl.create 20) in
-    analyse scope 0 args program
+    analyse scope 0 program
 
 let rec phase input =
     log_line 2 "- Context analysis";
     match input with
-    | Types node -> Types (analyse_context args node)
+    | Ast node -> Ast (analyse_context node)
     | _ -> raise (InvalidInput "context analysis")

phases/desug.ml

@@ -53,46 +53,45 @@ let rec var_init = function
                alloc)
 
     (* Split local variable initialisations in declaration and assignment *)
-    | FunDef (export, ret_type, name, params, Block body, ann) ->
-        let move_inits body =
-            let rec trav inits node = match node with
-                (* Translate scalar array initialisation to ArrayScalar node,
-                 * for easy replacement later on *)
-                | VarDec (Array _ as vtype, name, Some (BoolConst _  as v), ann) :: tl
-                | VarDec (Array _ as vtype, name, Some (FloatConst _ as v), ann) :: tl
-                | VarDec (Array _ as vtype, name, Some (IntConst _   as v), ann) :: tl ->
-                    let init = Some (ArrayInit (ArrayScalar v, vtype)) in
-                    trav inits (VarDec (vtype, name, init, ann) :: tl)
+    | FunDef (export, ret_type, name, params, body, ann) ->
+        let inits = ref [] in
+        let rec extract_inits = function
+            (* Translate scalar array initialisation to ArrayScalar node,
+                * for easy replacement later on *)
+            | VarDec (Array _ as vtype, name, Some (Const _   as v), ann) ->
+                let init = Some (ArrayInit (ArrayScalar v, vtype)) in
+                extract_inits (VarDec (vtype, name, init, ann))
 
-                (* Wrap ArrayConst in ArrayInit to pass dimensions *)
-                | VarDec (Array _ as vtype, name, Some (ArrayConst _ as v), ann) :: tl ->
-                    let init = Some (ArrayInit (v, vtype)) in
-                    trav inits (VarDec (vtype, name, init, ann) :: tl)
+            (* Wrap ArrayConst in ArrayInit to pass dimensions *)
+            | VarDec (Array _ as vtype, name, Some (ArrayConst _ as v), ann) ->
+                let init = Some (ArrayInit (v, vtype)) in
+                extract_inits (VarDec (vtype, name, init, ann))
 
-                | VarDec (ctype, name, init, ann) as dec :: tl ->
-                    (* array definition: create __allocate statement *)
-                    let alloc = match ctype with
-                        | Array (_, dims) -> [Allocate (name, dims, dec, ann)]
-                        | _ -> []
-                    in
-                    (* initialisation: create assign statement *)
-                    let add = match init with
-                        | Some value -> alloc @ [Assign (name, None, value, ann)]
-                        | None -> alloc
-                    in
-                    VarDec (ctype, name, None, ann) :: (trav (inits @ add) tl)
-
-                (* initialisations need to be placed after local functions *)
-                | (FunDef (_, _, _, _, _, _) as h) :: tl ->
-                    (var_init h) :: (trav inits tl)
+            | VarDec (ctype, name, init, ann) as dec ->
+                (* array definition: create __allocate statement *)
+                let alloc = match ctype with
+                    | Array (_, dims) -> [Allocate (name, dims, dec, ann)]
+                    | _ -> []
+                in
+                (* initialisation: create assign statement *)
+                let add = match init with
+                    | Some value -> alloc @ [Assign (name, None, value, ann)]
+                    | None -> alloc
+                in
+                inits := !inits @ add;
+                VarDec (ctype, name, None, ann)
 
-                (* rest of function body: recurse *)
-                | rest -> inits @ (List.map var_init rest)
-            in
-            flatten_blocks (trav [] body)
+            | node -> transform_children extract_inits node
+        in
+        let rec place_inits = function
+            (* initialisations need to be placed after local functions *)
+            | (LocalFuns _ as hd) :: tl -> hd :: !inits @ tl
+            | hd :: tl -> hd :: (place_inits tl)
+            | [] -> []
         in
         let params = flatten_blocks (List.map var_init params) in
-        FunDef (export, ret_type, name, params, Block (move_inits body), ann)
+        let body = flatten_blocks (place_inits (block_body (extract_inits body))) in
+        FunDef (export, ret_type, name, params, Block body, ann)
 
     | node -> transform_children var_init node
 
@@ -127,7 +126,7 @@ let for_to_while node =
             let vstop = Var (_stop, None, annof stop) in
             let vstep = Var (_step, None, annof step) in
             let cond = Cond (
-                Binop (Gt, vstep, IntConst (0, []), []),
+                Binop (Gt, vstep, Const (IntVal 0, []), []),
                 Binop (Lt, vi, vstop, []),
                 Binop (Gt, vi, vstop, []),
                 []
@@ -161,7 +160,7 @@ let rec array_init = function
             | dim :: rest ->
                 let counter = fresh_var "i" in
                 let body = Block [add_loop (indices @ [Var (counter, None, [])]) rest] in
-                For (counter, IntConst (0, []), dim, IntConst (1, []), body, [])
+                For (counter, Const (IntVal 0, []), dim, Const (IntVal 1, []), body, [])
         in
         add_loop [] dims
 
@@ -179,7 +178,7 @@ let rec array_init = function
             | ArrayConst (values, _) ->
                 make_assigns (depth + 1) 0 indices values
             | value when depth = ndims ->
-                let indices = List.map (fun i -> IntConst (i, [])) indices in
+                let indices = List.map (fun i -> Const (IntVal i, [])) indices in
                 [Assign (name, Some (List.rev indices), value, ann)]
             | node ->
                 let msg = sprintf
@@ -253,6 +252,6 @@ let rec array_dims = function
 let rec phase input =
     log_line 2 "- Desugaring";
     match input with
-    | Types node ->
-            Types (for_to_while (array_init (var_init (array_dims node))))
+    | Ast node ->
+            Ast (for_to_while (array_init (var_init (array_dims node))))
     | _ -> raise (InvalidInput "desugar")

phases/dim_reduce.ml

@@ -50,5 +50,5 @@ let rec simplify_decs = function
 let rec phase input =
     log_line 2 "- Array dimension reduction";
     match input with
-    | Types node -> Types (simplify_decs (dim_reduce 0 node))
+    | Ast node -> Ast (simplify_decs (dim_reduce 0 node))
     | _ -> raise (InvalidInput "dimension reduction")

phases/expand_dims.ml

@@ -38,5 +38,5 @@ let rec expand_dims = function
 let rec phase input =
     log_line 2 "- Expand array dimensions";
     match input with
-    | Types node -> Types (expand_dims node)
+    | Ast node -> Ast (expand_dims node)
     | _ -> raise (InvalidInput "expand dimensions")

phases/extern_vars.ml

@@ -126,8 +126,8 @@ let rec replace_vars scope depth = function
 let rec phase input =
     log_line 2 "- Create getters and setters for extern variables";
     match input with
-    | Types node ->
+    | Ast node ->
         let globals = Hashtbl.create 20 in
         let node = create_funcs globals node in
-        Types (replace_vars globals 0 node)
+        Ast (replace_vars globals 0 node)
     | _ -> raise (InvalidInput "extern vars")

phases/parse.ml

@@ -25,5 +25,5 @@ let phase input =
         let ast = parse_with_error lexbuf in
         (match ast with
             | None -> raise (CompileError "no syntax tree was constructed")
-            | Some node -> Types node)
+            | Some node -> Ast node)
     | _ -> raise (InvalidInput "parse")

phases/print.ml

@@ -4,6 +4,8 @@ open Stringify
 
 let tab = "    "
 
+let si = string_of_int
+
 let rec repeat s n = if n < 1 then "" else s ^ (repeat s (n - 1))
 let pad width s = s ^ (repeat " " (String.length s - width))
 let paddall width = List.map (pad width)
@@ -13,12 +15,12 @@ let type2str = function
     | Array (t, dims) -> ctype2str t ^ repeat "," (List.length dims - 1)
     | t -> ctype2str t
 
-let prefix node = match typeof node with
+let prefix = function
     | Bool _  -> "b"
     | Int _   -> "i"
     | Float _ -> "f"
-    | Array _ -> "a"
-    | _ -> raise InvalidNode
+    | Void  _ -> ""
+    | _       -> "a"
 
 let instr2str = function
     (* Global / directives *)
@@ -32,23 +34,39 @@ let instr2str = function
     | Import (name, ret_type, arg_types) ->
         let types = List.map type2str (ret_type :: arg_types) in
         ".import \"" ^ name ^ "\" " ^ (String.concat " " types)
-    | Const node ->
-        ".const  " ^ (node2str node)
+    | ConstDef node ->
+        ".const  " ^ (const2str node)
     | Global ctype ->
         ".global " ^ (type2str ctype)
 
+    (* Store *)
+    | StoreGlob (ctype, index) ->
+        tab ^ prefix ctype ^ "storeg " ^ si index
+    | StoreLoc (ctype, index) ->
+        tab ^ prefix ctype ^ "store  " ^ si index
+    | StoreRel (ctype, nesting, index) ->
+        tab ^ prefix ctype ^ "storen " ^ si nesting ^ " " ^ si index
+
     (* Load constant *)
     | LoadConst (ctype, index) ->
-        tab ^ type2str ctype ^ "loadc " ^ string_of_int index
-    | LoadImm (BoolConst (b, _)) ->
+        tab ^ prefix ctype ^ "loadc " ^ si index
+    | LoadImm (Const (BoolVal b, _)) ->
         tab ^ "bloadc_" ^ (if b then "t" else "f")
-    | LoadImm (IntConst (i, _)) when i < 0 ->
-        tab ^ "iloadc_m" ^ string_of_int (-i)
-    | LoadImm (IntConst (i, _)) ->
-        tab ^ "iloadc_" ^ string_of_int i
-    | LoadImm (FloatConst (i, _)) ->
-        tab ^ "floadc_" ^ string_of_int (int_of_float i)
+    | LoadImm (Const (IntVal i, _)) when i < 0 ->
+        tab ^ "iloadc_m" ^ si (-i)
+    | LoadImm (Const (IntVal i, _)) ->
+        tab ^ "iloadc_" ^ si i
+    | LoadImm (Const (FloatVal i, _)) ->
+        tab ^ "floadc_" ^ si (int_of_float i)
+
+    (* Control flow *)
+    | RtnEnter stack_len ->
+        tab ^ "esr " ^ si stack_len
+    | Ret ctype ->
+        tab ^ prefix ctype ^ "return"
 
+    | EmptyLine -> ""
+    | DummyInstr -> tab ^ "<dummy>"
     | _ -> tab ^ "<unknown instruction>"
 
 let rec print_assembly oc instrs =
@@ -69,15 +87,14 @@ let rec print_assembly oc instrs =
             trav tl
     in
     trav instrs;
-    List.iter output_line (List.rev !endbuf)
-    (*
-    let hasdot ins = if String.length ins > 0 && ins.[0] = '.' then 1 else 0 in
-    let cmp a b = compare (hasdot a) (hasdot b) in
-    List.sort cmp (trav instrs)
-    *)
+    if List.length !endbuf > 1 then (
+        output_line (instr2str (Comment ("globals:")));
+        let cmp a b = compare (String.sub a 0 8) (String.sub b 0 8) in
+        List.iter output_line (List.sort cmp (List.rev !endbuf))
+    ); ()
 
 let phase = function
-    | Types node as input ->
+    | Ast node as input ->
         if args.verbose >= 2 then (
             prerr_endline "--------------------------------------------------";
             prerr_endline (node2str node);
@@ -105,9 +122,6 @@ let phase = function
                 prerr_endline "--------------------------------------------------";
 
             print_assembly stdout instrs;
-
-            if args.verbose >= 2 then
-                prerr_endline "--------------------------------------------------"
         );
         input
 

phases/typecheck.ml

@@ -175,9 +175,9 @@ let rec typecheck node =
         DoWhile (check_trav Bool cond, typecheck body, ann)
 
     (* Constants *)
-    | BoolConst  (value, ann) -> BoolConst  (value, Type Bool  :: ann)
-    | IntConst   (value, ann) -> IntConst   (value, Type Int   :: ann)
-    | FloatConst (value, ann) -> FloatConst (value, Type Float :: ann)
+    | Const (BoolVal  value, ann) -> Const (BoolVal  value, Type Bool  :: ann)
+    | Const (IntVal   value, ann) -> Const (IntVal   value, Type Int   :: ann)
+    | Const (FloatVal value, ann) -> Const (FloatVal value, Type Float :: ann)
 
     (* Variables inherit the type of their declaration *)
     | VarUse (dec, None, ann) ->
@@ -213,5 +213,5 @@ let rec typecheck node =
 let rec phase input =
     log_line 2 "- Type checking";
     match input with
-    | Types node -> Types (typecheck node)
+    | Ast node -> Ast (typecheck node)
     | _ -> raise (InvalidInput "typecheck")

stringify.ml

@@ -5,11 +5,16 @@ let tab = "    "
 (* string -> string *)
 let indent = Str.global_replace (Str.regexp "^\\(.\\)") (tab ^ "\\1")
 
-(* Add a trailing zero to a float stringification
- * float -> string *)
-let float2string f = match string_of_float f with
-    | s when s.[String.length s - 1] = '.' -> s ^ "0"
-    | s -> s
+(* const -> string *)
+let const2str = function
+    | BoolVal  b -> string_of_bool b
+    | IntVal   i -> string_of_int i
+    | FloatVal f ->
+        (* Add a trailing zero to a float stringification *)
+        (match string_of_float f with
+        | s when s.[String.length s - 1] = '.' -> s ^ "0"
+        | s -> s
+        )
 
 (* Copied from util.ml to avoid circular dependency *)
 let nameof = function
@@ -102,19 +107,24 @@ and node2str node =
         "do " ^ str body ^ " while (" ^ str cond ^ ");"
     | For (counter, start, stop, step, body, _) ->
         let step = match step with
-            | IntConst (1, _) -> ""
+            | Const (IntVal 1, _) -> ""
             | value -> ", " ^ str value
         in
         let range = str start ^ ", " ^ str stop ^ step in
         "for (int " ^ counter ^ " = " ^ range ^ ") " ^ str body
     | Allocate (name, dims, _, _) ->
         name ^ " := <allocate>(" ^ concat ", " dims ^ ");"
-    | Block body -> "{\n" ^ indent (concat "\n" body) ^ "\n}"
+    | Block body ->
+        let rec append = function
+            | [] -> ""
+            | [last] -> last
+            | "" :: tl -> append tl
+            | hd :: tl -> hd ^ "\n" ^ append tl
+        in
+        "{\n" ^ indent (append (List.map str body)) ^ "\n}"
 
     (* Expressions *)
-    | BoolConst (b, _) -> string_of_bool b
-    | IntConst (i, _) -> string_of_int i
-    | FloatConst (f, _) -> float2string f
+    | Const (c, _) -> const2str c
     | ArrayConst (dims, _) -> "[" ^ concat ", " dims ^ "]"
     | Var (v, None, _) -> v
     | Var (name, Some dims, _) -> name ^ (str (ArrayConst (dims, [])))
@@ -147,11 +157,13 @@ and node2str node =
     | FunUse (dec, args, _) ->
         node2str (FunCall (nameof dec, args, []))
     | Dim (name, _) -> name
-
     | ArrayScalar node
     | ArrayInit (node, _)
     | Arg node -> str node
 
+    | VarDecs nodes
+    | LocalFuns nodes -> concat "\n" nodes
+
     | DummyNode -> "<dummy>"
 
 (* ctype list -> string *)

stringify.mli

@@ -1,7 +1,9 @@
-val op2str : Types.operator -> string
+val const2str : Types.const -> string
 
-val node2str : Types.node -> string
+val op2str    : Types.operator -> string
 
-val type2str : Types.ctype -> string
+val node2str  : Types.node -> string
+
+val type2str  : Types.ctype -> string
 
 val types2str : Types.ctype list -> string

test/old/array_dim.cvc

@@ -4,9 +4,8 @@ int dim() {
     dims = dims + 1;
     return 10;
 }
-/*
+
 void foo() {
    int[10, dim()] arr;
    arr[1, 0] = 1;
 }
-*/

types.ml

@@ -5,6 +5,7 @@ type operator = Neg | Not
               | Add | Sub | Mul | Div | Mod
               | Eq | Ne | Lt | Le | Gt | Ge
               | And | Or
+type const = BoolVal of bool | IntVal of int | FloatVal of float
 type ctype = Void | Bool | Int | Float
            | Array of ctype * node list
            | ArrayDepth of ctype * int  (* TODO: remove? *)
@@ -12,6 +13,7 @@ type ctype = Void | Bool | Int | Float
 and annotation =
     | Loc of location
     | Depth of int
+    | Index of int
     | Type of ctype
 and ann = annotation list
 and node =
@@ -31,6 +33,9 @@ and node =
     | Dim of string * ann
       (* dimension name in array Param *)
 
+    | VarDecs of node list
+    | LocalFuns of node list
+
     (* Statements *)
     | VarDec of ctype * string * node option * ann
       (* type, name, initialisation? *)
@@ -49,9 +54,7 @@ and node =
     | DoWhile of node * node * ann           (* cond, body *)
 
     (* Expressions *)
-    | BoolConst of bool * ann                  (* bool value *)
-    | IntConst of int * ann                    (* int value *)
-    | FloatConst of float * ann                (* float value *)
+    | Const of const * ann                     (* bool|int|float value *)
     | ArrayConst of node list * ann            (* [<exprs>] *)
     | Var of string * node list option * ann   (* <name> [<indices>]? *)
     | Monop of operator * node * ann           (* op, operand *)
@@ -82,22 +85,30 @@ type instr =
     (* .import "<name>" <ret_type> [ <arg_type>; ... ] *)
     | Import of string * ctype * ctype list
     (* .const <value> *)
-    | Const of node
+    | ConstDef of const
     (* .global <type> *)
     | Global of ctype
 
-    (* [ifb]loadg G *)
-    | LoadGlob of node
-    (* [ifb]loadc C *)
-    | LoadConst of ctype * int
-    (* [ifb]load_[01tf] <value> *)
-    | LoadImm of node
+    | StoreGlob of ctype * int          (* [ifba]storeg G *)
+    | StoreLoc of ctype * int           (* [ifba]store L *)
+    | StoreRel of ctype * int * int     (* [ifba]storen N L *)
+
+    | LoadGlob of node                  (* [ifb]loadg G *)
+    | LoadConst of ctype * int          (* [ifb]loadc C *)
+    | LoadImm of node                   (* [ifb]load_[01tf] <value> *)
+
+    (* Control flow *)
+    | RtnEnter of int
+    | RtnInit
+    | RtnJmp
+    | Ret of ctype
 
     (* Instructions *)
-    (* i(inc|dec) L C *)
-    | Inc of int * int
-    (* i(inc|dec)_1 C *)
-    | IncOne of int
+    | Inc of int * int                  (* i(inc|dec) L C *)
+    | IncOne of int                     (* i(inc|dec)_1 C *)
+
+    | EmptyLine
+    | DummyInstr
 
 (* Container for command-line arguments *)
 type args_record = {
@@ -112,8 +123,8 @@ let verbosity_default = 2  (* TODO: set to 1 when done with debugging *)
 let verbosity_debug   = 3
 
 (* Commandline args are stored in a global struct
- * (yes, it IS dirty, but I'd rather have this than having to pass [args] around
- * everywhere) *)
+ * (yes, it IS dirty, but I don't know how to do this without passin [args] to
+ * every function) *)
 let args = {
     infile  = None;
     outfile = None;
@@ -125,7 +136,7 @@ let args = {
 type intermediate =
     | Empty
     | FileContent of string * string
-    | Types of node
+    | Ast of node
     | Assembly of instr list
 
 (* exceptions *)

util.ml

@@ -44,18 +44,18 @@ let loc_from_lexpos pstart pend =
 let rec flatten_blocks lst =
     let flatten = flatten_blocks in
     let trav = function
-        | FunDef (export, ret_type, name, params, Block body, loc) ->
-            FunDef (export, ret_type, name, flatten params, Block (flatten body), loc)
-        | If (cond, Block body, loc) ->
-            If (cond, Block (flatten body), loc)
-        | IfElse (cond, Block tbody, Block fbody, loc) ->
-            IfElse (cond, Block (flatten tbody), Block (flatten fbody), loc)
-        | While (cond, Block body, loc) ->
-            While (cond, Block (flatten body), loc)
-        | DoWhile (cond, Block body, loc) ->
-            DoWhile (cond, Block (flatten body), loc)
-        | For (counter, start, stop, step, Block body, loc) ->
-            For (counter, start, stop, step, Block (flatten body), loc)
+        | FunDef (export, ret_type, name, params, Block body, ann) ->
+            FunDef (export, ret_type, name, flatten params, Block (flatten body), ann)
+        | If (cond, Block body, ann) ->
+            If (cond, Block (flatten body), ann)
+        | IfElse (cond, Block tbody, Block fbody, ann) ->
+            IfElse (cond, Block (flatten tbody), Block (flatten fbody), ann)
+        | While (cond, Block body, ann) ->
+            While (cond, Block (flatten body), ann)
+        | DoWhile (cond, Block body, ann) ->
+            DoWhile (cond, Block (flatten body), ann)
+        | For (counter, start, stop, step, Block body, ann) ->
+            For (counter, start, stop, step, Block (flatten body), ann)
         | node -> node
     in
     match lst with
@@ -69,56 +69,56 @@ let rec flatten_blocks lst =
 let transform_children trav node =
     let trav_all nodes = List.map trav nodes in
     match node with
-    | Program (decls, loc) ->
-        Program (flatten_blocks (trav_all decls), loc)
-    | FunDec (ret_type, name, params, loc) ->
-        FunDec (ret_type, name, trav_all params, loc)
-    | FunDef (export, ret_type, name, params, body, loc) ->
-        FunDef (export, ret_type, name, trav_all params, trav body, loc)
-    | GlobalDec (ctype, name, loc) ->
-        GlobalDec (ctype, name, loc)
-    | GlobalDef (export, ctype, name, Some init, loc) ->
-        GlobalDef (export, ctype, name, Some (trav init), loc)
-
-    | VarDec (ctype, name, Some init, loc) ->
-        VarDec (ctype, name, Some (trav init), loc)
-    | Assign (name, None, value, loc) ->
-        Assign (name, None, trav value, loc)
-    | Assign (name, Some dims, value, loc) ->
-        Assign (name, Some (trav_all dims), trav value, loc)
-    | VarLet (dec, None, value, loc) ->
-        VarLet (dec, None, trav value, loc)
-    | VarLet (dec, Some dims, value, loc) ->
-        VarLet (dec, Some (trav_all dims), trav value, loc)
-    | Return (value, loc) ->
-        Return (trav value, loc)
-    | If (cond, body, loc) ->
-        If (trav cond, trav body, loc)
-    | IfElse (cond, true_body, false_body, loc) ->
-        IfElse (trav cond, trav true_body, trav false_body, loc)
-    | While (cond, body, loc) ->
-        While (trav cond, trav body, loc)
-    | DoWhile (cond, body, loc) ->
-        DoWhile (trav cond, trav body, loc)
-    | For (counter, start, stop, step, body, loc) ->
-        For (counter, trav start, trav stop, trav step, trav body, loc)
-    | Allocate (name, dims, dec, loc) ->
-        Allocate (name, trav_all dims, dec, loc)
+    | Program (decls, ann) ->
+        Program (flatten_blocks (trav_all decls), ann)
+    | FunDec (ret_type, name, params, ann) ->
+        FunDec (ret_type, name, trav_all params, ann)
+    | FunDef (export, ret_type, name, params, body, ann) ->
+        FunDef (export, ret_type, name, trav_all params, trav body, ann)
+    | GlobalDec (ctype, name, ann) ->
+        GlobalDec (ctype, name, ann)
+    | GlobalDef (export, ctype, name, Some init, ann) ->
+        GlobalDef (export, ctype, name, Some (trav init), ann)
+
+    | VarDec (ctype, name, Some init, ann) ->
+        VarDec (ctype, name, Some (trav init), ann)
+    | Assign (name, None, value, ann) ->
+        Assign (name, None, trav value, ann)
+    | Assign (name, Some dims, value, ann) ->
+        Assign (name, Some (trav_all dims), trav value, ann)
+    | VarLet (dec, None, value, ann) ->
+        VarLet (dec, None, trav value, ann)
+    | VarLet (dec, Some dims, value, ann) ->
+        VarLet (dec, Some (trav_all dims), trav value, ann)
+    | Return (value, ann) ->
+        Return (trav value, ann)
+    | If (cond, body, ann) ->
+        If (trav cond, trav body, ann)
+    | IfElse (cond, true_body, false_body, ann) ->
+        IfElse (trav cond, trav true_body, trav false_body, ann)
+    | While (cond, body, ann) ->
+        While (trav cond, trav body, ann)
+    | DoWhile (cond, body, ann) ->
+        DoWhile (trav cond, trav body, ann)
+    | For (counter, start, stop, step, body, ann) ->
+        For (counter, trav start, trav stop, trav step, trav body, ann)
+    | Allocate (name, dims, dec, ann) ->
+        Allocate (name, trav_all dims, dec, ann)
     | Expr value ->
         Expr (trav value)
     | Block (body) ->
         Block (trav_all body)
 
-    | Monop (op, value, loc) ->
-        Monop (op, trav value, loc)
-    | Binop (op, left, right, loc) ->
-        Binop (op, trav left, trav right, loc)
-    | Cond (cond, true_expr, false_expr, loc) ->
-        Cond (trav cond, trav true_expr, trav false_expr, loc)
-    | TypeCast (ctype, value, loc) ->
-        TypeCast (ctype, trav value, loc)
-    | FunCall (name, args, loc) ->
-        FunCall (name, trav_all args, loc)
+    | Monop (op, value, ann) ->
+        Monop (op, trav value, ann)
+    | Binop (op, left, right, ann) ->
+        Binop (op, trav left, trav right, ann)
+    | Cond (cond, true_expr, false_expr, ann) ->
+        Cond (trav cond, trav true_expr, trav false_expr, ann)
+    | TypeCast (ctype, value, ann) ->
+        TypeCast (ctype, trav value, ann)
+    | FunCall (name, args, ann) ->
+        FunCall (name, trav_all args, ann)
     | Arg value ->
         Arg (trav value)
 
@@ -133,9 +133,78 @@ let transform_children trav node =
     | FunUse (dec, params, ann) ->
         FunUse (dec, trav_all params, ann)
 
+    | VarDecs decs ->
+        VarDecs (trav_all decs)
+    | LocalFuns funs ->
+        LocalFuns (trav_all funs)
+
     | _ -> node
 
- let rec annof = function
+let annotate a = function
+    | Program (decls, ann) ->
+        Program (decls, a :: ann)
+    | FunDec (ret_type, name, params, ann) ->
+        FunDec (ret_type, name, params, a :: ann)
+    | FunDef (export, ret_type, name, params, body, ann) ->
+        FunDef (export, ret_type, name, params, body, a :: ann)
+    | GlobalDec (ctype, name, ann) ->
+        GlobalDec (ctype, name, a :: ann)
+    | GlobalDef (export, ctype, name, init, ann) ->
+        GlobalDef (export, ctype, name, init, a :: ann)
+    | VarDec (ctype, name, init, ann) ->
+        VarDec (ctype, name, init, a :: ann)
+    | Assign (name, dims, value, ann) ->
+        Assign (name, dims, value, a :: ann)
+    | VarLet (dec, dims, value, ann) ->
+        VarLet (dec, dims, value, a :: ann)
+    | Return (value, ann) ->
+        Return (value, a :: ann)
+    | If (cond, body, ann) ->
+        If (cond, body, a :: ann)
+    | IfElse (cond, true_body, false_body, ann) ->
+        IfElse (cond, true_body, false_body, a :: ann)
+    | While (cond, body, ann) ->
+        While (cond, body, a :: ann)
+    | DoWhile (cond, body, ann) ->
+        DoWhile (cond, body, a :: ann)
+    | For (counter, start, stop, step, body, ann) ->
+        For (counter, start, stop, step, body, a :: ann)
+    | Allocate (name, dims, dec, ann) ->
+        Allocate (name, dims, dec, a :: ann)
+    | Monop (op, value, ann) ->
+        Monop (op, value, a :: ann)
+    | Binop (op, left, right, ann) ->
+        Binop (op, left, right, a :: ann)
+    | Cond (cond, true_expr, false_expr, ann) ->
+        Cond (cond, true_expr, false_expr, a :: ann)
+    | TypeCast (ctype, value, ann) ->
+        TypeCast (ctype, value, a :: ann)
+    | FunCall (name, args, ann) ->
+        FunCall (name, args, a :: ann)
+    | Arg value ->
+        Arg (value)
+    | Var (dec, dims, ann) ->
+        Var (dec, dims, a :: ann)
+    | VarUse (dec, dims, ann) ->
+        VarUse (dec, dims, a :: ann)
+    | FunUse (dec, params, ann) ->
+        FunUse (dec, params, a :: ann)
+    | Const (BoolVal value, ann) ->
+        Const (BoolVal value, a :: ann)
+    | Const (IntVal value, ann) ->
+        Const (IntVal value, a :: ann)
+    | Const (FloatVal value, ann) ->
+        Const (FloatVal value, a :: ann)
+    | ArrayConst (value, ann) ->
+        ArrayConst (value, a :: ann)
+    | Param (ctype, name, ann) ->
+        Param (ctype, name, a :: ann)
+    | Dim (name, ann) ->
+        Dim (name, a :: ann)
+
+    | _ -> raise InvalidNode
+
+let rec annof = function
     | Program (_, ann)
     | Param (_, _, ann)
     | Dim (_, ann)
@@ -153,9 +222,9 @@ let transform_children trav node =
     | DoWhile (_, _, ann)
     | For (_, _, _, _, _, ann)
     | Allocate (_, _, _, ann)
-    | BoolConst (_, ann)
-    | IntConst (_, ann)
-    | FloatConst (_, ann)
+    | Const (BoolVal _, ann)
+    | Const (IntVal _, ann)
+    | Const (FloatVal _, ann)
     | ArrayConst (_, ann)
     | Var (_, _, ann)
     | Monop (_, _, ann)
@@ -173,20 +242,33 @@ let transform_children trav node =
 
     | _ -> raise InvalidNode
 
- let locof node =
+let locof node =
      let rec trav = function
          | []            -> noloc
          | Loc loc :: tl -> loc
          | _ :: tl       -> trav tl
      in trav (annof node)
 
- let rec depthof node =
+let rec depthof node =
      let rec trav = function
-         | []                -> raise InvalidNode
+         | [] ->
+            prerr_string "cannot get depth for: ";
+            prt_node node;
+            raise InvalidNode
          | Depth depth :: tl -> depth
          | _ :: tl           -> trav tl
      in trav (annof node)
 
+let rec indexof node =
+     let rec trav = function
+         | [] ->
+            prerr_string "cannot get index for: ";
+            prt_node node;
+            raise InvalidNode
+         | Index index :: tl -> index
+         | _ :: tl           -> trav tl
+     in trav (annof node)
+
 let typeof = function
     (* Some nodes have their type as property *)
     | VarDec (ctype, _, _, _)
@@ -201,7 +283,10 @@ let typeof = function
     (* Other nodes must be annotated during typechecking *)
     | node ->
         let rec trav = function
-            | []           -> prerr_string "cannot get type for: "; prt_node node; raise InvalidNode
+         | [] ->
+            prerr_string "cannot get type for: ";
+            prt_node node;
+            raise InvalidNode
             | Type t :: tl -> t
             | _ :: tl      -> trav tl
         in trav (annof node)
@@ -210,7 +295,7 @@ let prerr_loc (fname, ystart, yend, xstart, xend) =
     let file = open_in fname in
 
     (* skip lines until the first matched line *)
-    for i = 1 to ystart - 1 do input_line file done;
+    for i = 1 to ystart - 1 do let _ = input_line file in () done;
 
     (* for each line in `loc`, print the source line with an underline *)
     for l = ystart to yend do

util.mli

@@ -18,12 +18,16 @@ val loc_from_lexpos : Lexing.position -> Lexing.position -> Types.location
 (* Default transformation traversal for AST nodes *)
 val transform_children : (Types.node -> Types.node) -> Types.node -> Types.node
 
+(* Add a single annotation to a node (no traversal) *)
+val annotate : Types.annotation -> Types.node -> Types.node
+
 (*val visit_children : (Types.node -> unit) -> Types.node -> unit*)
 
 (* Extract annotation from node *)
 val annof   : Types.node -> Types.annotation list
 val locof   : Types.node -> Types.location
 val depthof : Types.node -> int
+val indexof : Types.node -> int
 val typeof  : Types.node -> Types.ctype
 
 (* Print file location to stderr *)