Go Compiler - Part 2

This is a continuation from last post.

After building the Ast what can we do next? We can check the types and compile the program.

This will be the expected:

...
let code_typed =
  [
    Func
      ( "add",
        [ ("a", TInt); ("b", TInt) ],
        TInt,
        [ Return (Op (Var "a", "+", Var "b")) ] );
    Func
      ( "main",
        [],
        TVoid,
        [
          Let ("a", Some TInt, Const 10);
          Let ("b", Some TInt, Const 25);
          Let ("_", Some TInt, Apply ("add", [ Var "a"; Var "b" ]));
        ] );
  ]
... 

Let's start by creating the type checker! First, let's create a new file typechecker.ml and define a function that takes our ast list and returns another ast list (but this time checked).

val check : typ VMap.t -> typ list VMap.t -> stmt list -> stmt list

First, let's use the Map data structure. It works identically to the HashTable in OCaml, but it's immutable, meaning that when we change the Map.v we need to store it again if we want to persist it. It is good to score the scope of variables. The VMap is the Map implementation we are using, uses a string as a key and stores the typ or typ list value (depending on the context).

module VMap = Map.Make (String)

let ctx = VMap.empty
let fun_ctx = VMap.empty

This will create a context to store functions (fun_ctx) and another one to store the variables (ctx).

Let's convert the Ast expressions to types.

let rec typ_of_ex ctx fun_ctx : Ast.expr -> Ast.typ = function
  | Const _ -> TInt
  | Op _ -> TInt
  | Var s -> VMap.find s ctx
  | Apply (s, exprs) ->
      let t_l = VMap.find s fun_ctx in
      let typs = List.map (typ_of_ex ctx fun_ctx) exprs in
      let rec h t1 t2 =
        match (t1, t2) with
        | [], [ t2 ] -> t2
        | t1 :: tl1, t2 :: tl2 when t1 = t2 -> h tl1 tl2
        | _ ->
            failwith "Not the right number of arguments or the type is invalid!"
      in
      h typs t_l;;

This function with take any expression and convert it to a type. Const and Operators are easy, as we only support ints, for now at least. With variables we have to find them in the context. The application is the hardest one, as we have to check the type of each expression and compare the types of each expression with the type required by the function.

This is getting the types of the function and the type of the expressions (one by one).

	...
	let t_l = VMap.find s fun_ctx in
	let typs = List.map (typ_of_ex ctx fun_ctx) exprs in
	...

This function compare both of the types, and returns the return type.

	...
	let rec h t1 t2 =
		match (t1, t2) with
		| [], [ t2 ] -> t2
		| t1 :: tl1, t2 :: tl2 when t1 = t2 -> h tl1 tl2
		| _ ->
			failwith "Not the right number of arguments or the type is invalid!"
		in
	h typs t_l;;

We will need a function to check the types:

let check_types (t1 : Ast.typ option) (t2 : Ast.typ) : Ast.typ =
  match t1 with
  | Some t1 -> if t1 = t2 then t1 else failwith "Types should be equal"
  | None -> t2

And finally, we can implement the check function

let rec check (ctx : typ VMap.t) (fun_ctx : typ list VMap.t) :
    Ast.stmt list -> Ast.stmt list = function
  | [] -> []
  | (Return _ as ret) :: tl -> ret :: check ctx fun_ctx tl
...

Note: OCaml can infer the types, we could write this as:

let rec check ctx fun_ctx = function
	...

First, we handle the empty case, and then the Return statement (we keep it equal).

	| Let (n, t_op, e) :: tl ->
	  let t = typ_of_ex ctx fun_ctx e in
	  let t = check_types t_op t in
	  let ctx = VMap.add n t ctx in
	  Let (n, Some t, e) :: check ctx fun_ctx tl

In the Let case, we will grab the type from the expression and compare it to the type we wrote in the source code (if any), then we add that variable to the context, we create a new let but this time with a type, and finally we continue iterating.

	...
	| Func (n, params, t_op, stmts) :: tl ->
      let rec stmts_typ (stmts : Ast.stmt list) (acc : Ast.stmt list * typ) =
        match stmts with
        | [] -> (fst acc, TVoid)
        | [ (Return e as ret) ] ->
            ( (check ctx fun_ctx [ ret ] |> List.hd) :: fst acc,
              typ_of_ex ctx fun_ctx e )
        | stmt :: tl ->
            stmts_typ tl
              ((check ctx fun_ctx [ stmt ] |> List.hd) :: fst acc, TVoid)
      in
      let stmts, t = stmts_typ stmts ([], TVoid) in
      let stmts = List.rev stmts in
      let t = check_types (Some t_op) t in
      let rec mount_t (params : (string * Ast.typ) list) =
        match params with [] -> [ t ] | (_, t) :: tl -> t :: mount_t tl
      in
      let ctx =
        List.fold_left (fun c (nn, typ) -> VMap.add nn typ c) ctx params
      in
      let t_mount = mount_t params in
      let fun_ctx = VMap.add n t_mount fun_ctx in
      Func (n, params, t, stmts) :: check ctx fun_ctx tl

For the Func case, first, we try to find the return type from the list of statements that will execute inside the function. After that, we check that type with the return type we wrote on the source code. Then, we add all params names to the variable context, then we group all types (the params type and the return type) into an array to know what to expect when calling the function. Finally, we return a new Func with the correct types.

This code is expected to run without any error.

...
let () =
  let lex = Lexing.from_channel (open_in "main.go") in
  let typed_code =
    Parser.code Lexer.token lex
    |> Typechecker.check Typechecker.ctx Typechecker.fun_ctx
  in
  assert (code_typed = typed_code)

Llvm Compilation... To be continued.