compile.ml

open Format
open X86_64
open Ast

let debug = ref false

let preamble = "
P_test: # argument in %rdi
      pushq   %rbp
      movq    %rsp, %rbp
      movq    (%rdi), %rax
      cmpq    $0, %rax
      je      E_test
      movq    8(%rdi), %rax
E_test:
      movq    %rbp, %rsp
      popq    %rbp
      ret
P_print_None:
      pushq   %rbp
      movq    %rsp, %rbp
      andq    $-16, %rsp
      movq    $S_message_None, %rdi
      xorq    %rax, %rax
      call    printf
      movq    %rbp, %rsp
      popq    %rbp
      ret
P_print_string:
      pushq   %rbp
      movq    %rsp, %rbp
      andq    $-16, %rsp
      movq    %rdi, %rsi
      movq    $S_message_string, %rdi
      xorq    %rax, %rax
      call    printf
      movq    %rbp, %rsp
      popq    %rbp
      ret
P_print_int:
      pushq   %rbp
      movq    %rsp, %rbp
      andq    $-16, %rsp
      movq    %rdi, %rsi
      movq    $S_message_int, %rdi
      xorq    %rax, %rax
      call    printf
      movq    %rbp, %rsp
      popq    %rbp
      ret

P_print_list:
    pushq   %rbp
    movq    %rsp, %rbp
    subq    $24, %rsp           
    pushq   %rbx             
    pushq   %r12
    pushq   %r13
    movq    %rdi, %rbx        
    movq    8(%rbx), %r12       
    xorq    %r13, %r13    
    xorq    %rax, %rax
    movq    $S_lb, %rdi
    call    printf
    testq   %r12, %r12
    je      PL_done

PL_loop:
    movq    16(%rbx,%r13,8), %rdi
    xorq    %rax, %rax
    call    P_print
    incq    %r13                   
    cmpq    %r13, %r12              
    je      PL_done
    xorq    %rax, %rax
    movq    $S_comma_space, %rdi
    call    printf
    jmp     PL_loop

PL_done:
    movq    $S_rb, %rdi
    call    printf
    popq    %r13                
    popq    %r12
    popq    %rbx
    leave
    ret

P_print:
      pushq   %rbp
      movq    %rsp, %rbp
      andq    $-16, %rsp
      movq    (%rdi), %rax      
      cmpq    $0, %rax           
      je      L0
      cmpq    $1, %rax           # Boolean?
      je      L1
      cmpq    $2, %rax           # Int?
      je      L2
      cmpq    $3, %rax           # String?
      je      L3
      cmpq    $4, %rax           # List?
      je      L4

L0:
      call    P_print_None
      jmp     E_print

L1:
      movq    8(%rdi), %rsi      
      cmpq    $0, %rsi
      je      L1_false
      movq    $S_message_True, %rdi
      call    printf
      jmp     E_print

L1_false:
      movq    $S_message_False, %rdi
      call    printf
      jmp     E_print

L2:
      movq    8(%rdi), %rdi
      call    P_print_int
      jmp     E_print

L3:
      leaq    16(%rdi), %rdi
      call    P_print_string
      jmp     E_print

L4:
      call    P_print_list
      jmp     E_print

E_print:
      movq    %rbp, %rsp
      popq    %rbp
      ret

P_alloc_int:
      pushq   %rbp
      movq    %rsp, %rbp
      pushq   %rdi        # the integer n to allocate is stored in %rdi
                          # and now we push it into the stack;
      andq    $-16, %rsp  # stack alignment;
      movq    $16, %rdi   # how many bytes you want to allocate;
      call    malloc      # the new allocated address is in %rax,
                          # which is a 16 bytes = 2 * 8 = 2 * 64 bits
                          # segment;
      movq    $2, (%rax)  # put the tag of an integer block
                          # in the address pointed by %rax;
      movq    -8(%rbp), %rdi # get back the value of n, which is now on
                             # the stack, at address %rbp - 8 bytes;
      movq    %rdi, 8(%rax)  # put the value of n at address
                             # %rax + 8 bytes;
      #### Now, we have the following, contiguous block allocated:
      ####
      ####        +---------+---------+
      ####        |    2    |    n    |
      ####        +---------+---------+
      ####        | 8 bytes | 8 bytes |
      ####        +---------+---------+
      ####
      movq    %rbp, %rsp
      popq    %rbp
      ret                 # the result is in %rax
P_alloc_list:
      pushq   %rbp
      movq    %rsp, %rbp
      pushq   %rdi        # the length of the list
      andq    $-32, %rsp  # stack alignment;
      shl     $3, %rdi    # 8 * %rdi
      addq    $16, %rdi   # %rdi = 16 + 8 * length of the list
      call    malloc
      movq    $4, (%rax)
      movq    -8(%rbp), %rdi
      movq    %rdi, 8(%rax)
      movq    %rbp, %rsp
      popq    %rbp 
      ret                 # the result is in %rax
P_print_newline:
      pushq   %rbp
      movq    %rsp, %rbp
      andq    $-16, %rsp
      movq    $S_newline, %rdi
      xorq    %rax, %rax
      call    printf
      movq    %rbp, %rsp
      popq    %rbp
      ret

P_add:
    pushq   %rbp
    movq    %rsp, %rbp

    # test for two strings?
    movq    (%rdi), %rax      # lhs.tag
    cmpq    $3, %rax
    jne     Ladd_int
    movq    (%rsi), %rax      # rhs.tag
    cmpq    $3, %rax
    jne     Ladd_int

    # both are strings → concat
    movq    %rdi, %rdi
    movq    %rsi, %rsi
    call    P_concat
    jmp     Ladd_done

Ladd_int:
    movq    %rdi, %rdi       
    movq    %rsi, %rsi
    call    P_add_int

Ladd_done:
    popq    %rbp
    ret

P_add_int:
    pushq   %rbp
    movq    %rsp, %rbp
    movq    (%rdi), %rax        
    cmpq    $2, %rax
    jne      type_error
    movq    8(%rdi), %rax      
    movq    (%rsi), %rcx
    cmpq    $2, %rcx
    jne      type_error
    movq    8(%rsi), %rcx      
    addq    %rcx, %rax         
    movq    %rax, %rdi
    call    P_alloc_int        
    leave
    ret


P_concat:
    pushq   %rbp
    movq    %rsp, %rbp

    pushq   %rsi
    pushq   %rdi

    # concat length
    movq    8(%rdi), %rdx
    addq    8(%rsi), %rdx

    # string malloc
    pushq   %rdx
    incq    %rdx
    movq    %rdx, %rdi
    addq    $16, %rdi
    call    malloc
    movq    $3, (%rax)
    popq    %rcx
    movq    %rcx, 8(%rax)

    # string concat
    popq    %rdi
    pushq   %rax
    leaq    16(%rdi), %rsi
    leaq    16(%rax), %rdi
    call    strcat
    popq    %rax

    popq    %rsi
    pushq   %rax
    leaq    16(%rsi), %rsi
    leaq    16(%rax), %rdi
    call    strcat
    popq    %rax

    popq    %rbp
    ret

P_sub_int:
    pushq   %rbp
    movq    %rsp, %rbp
    movq    (%rdi), %rax
    cmpq    $2, %rax
    jne      type_error
    movq    8(%rdi), %rax
    movq    (%rsi), %rcx
    cmpq    $2, %rcx
    jne      type_error
    movq    8(%rsi), %rcx
    subq    %rcx, %rax         
    movq    %rax, %rdi
    call    P_alloc_int
    leave
    ret

P_div_int:
    pushq   %rbp
    movq    %rsp, %rbp
    movq    (%rdi), %rax
    cmpq    $2, %rax
    jne      type_error
    movq    8(%rdi), %rax       
    movq    (%rsi), %rcx
    cmpq    $2, %rcx
    jne      type_error
    movq    8(%rsi), %rcx       
    testq   %rcx, %rcx
    je       div_zero
    cqto
    idivq   %rcx              
    testq   %rdx, %rdx
    jz       div_ok             
    movq    8(%rdi), %r8       
    xorq    8(%rsi), %r8     
    testq   %r8, %r8
    jns      div_ok            
    decq    %rax              

div_ok:
    movq    %rax, %rdi
    call    P_alloc_int
    leave
    ret

div_zero:
    movq    $S_error_div_zero, %rdi
    call    P_runtime_error


P_mod_int:
    pushq   %rbp
    movq    %rsp, %rbp
    movq    (%rdi), %rax
    cmpq    $2, %rax
    jne      type_error
    movq    8(%rdi), %rax
    movq    (%rsi), %rcx
    cmpq    $2, %rcx
    jne      type_error
    movq    8(%rsi), %rcx
    testq   %rcx, %rcx
    je       mod_zero
    cqto
    idivq   %rcx               
    testq   %rdx, %rdx
    jz       mod_ok
    movq    8(%rdi), %r8
    xorq    8(%rsi), %r8
    testq   %r8, %r8
    jns      mod_ok
    addq    %rcx, %rdx         

mod_ok:
    movq    %rdx, %rdi
    call    P_alloc_int
    leave
    ret
mod_zero:
    movq    $S_error_div_zero, %rdi
    call    P_runtime_error

compare_values:
    push %rbp
    mov %rsp, %rbp
    movq (%rdi), %rdx
    movq (%rsi), %rcx
    cmpq %rdx, %rcx
    jne  type_mismatch
    movq %rdx, %r8
    movq 8(%rdi), %rdx
    movq 8(%rsi), %rcx
    movq %rdx, %rax
    subq %rcx, %rax
    cmpq $0, %rax
    jne  done
    cmpq $0, %rdx
    je  done
    cmpq $3, %r8
    je  compare_strings
    cmpq $4, %r8
    je  compare_lists
    jmp  done

compare_strings:
    leaq 16(%rdi), %rdi
    leaq 16(%rsi), %rsi
    call strcmp           
    jmp  done

compare_lists:
    movq $0, %r8


list_loop:
    movq 16(%rdi, %r8, 8), %rcx
    movq 16(%rsi, %r8, 8), %r9
    pushq %rdx
    pushq %r8
    pushq %rdi
    pushq %rsi
    movq %rcx, %rdi
    movq %r9, %rsi
    call compare_values
    popq %rsi
    popq %rdi
    popq %r8
    popq %rdx
    cmpq $0, %rax
    jne  done
    incq %r8
    cmpq %rdx, %r8
    jne  list_loop
    jmp  done
    
type_mismatch:
    cmpq  $0, %r10
    je    tm_false     
    cmpq  $1, %r10
    je    tm_true      
    movq  $S_error_compare_type, %rdi
    call  P_runtime_error

tm_false:
    movq  $1, %rax      # non-zero, so '==?' will see rax!=0 → no je
    jmp   done

tm_true:
    movq  $1, %rax      # non-zero, so '!=?' will see rax!=0 → take jne
    jmp   done

P_runtime_error:
    pushq   %rbp
    movq    %rsp, %rbp
    andq    $-16, %rsp
    movq    %rdi, %rsi
    leaq    S_error_fmt(%rip), %rdi
    xorq    %rax, %rax
    call    printf
    movq    $1, %rdi
    call    exit

P_index:
    pushq   %rbp
    movq    %rsp, %rbp
    movq    (%rdi), %rax       
    movq    8(%rdi), %rdx      
    cmpq    $0, %rsi
    jge      idx_nonneg
    addq    %rdx, %rsi
idx_nonneg:
    cmpq    $0, %rsi
    jl       PERR_index
    cmpq    %rdx, %rsi
    jge      PERR_index
    cmpq    $4, %rax           
    je       do_list
    cmpq    $3, %rax           
    je       do_string
    jmp      PERR_type

  do_list:
    shlq    $3, %rsi          
    addq    $16, %rsi          
    addq    %rsi, %rdi          
    movq    (%rdi), %rdi       
    jmp      P_done

  do_string:
    movb    16(%rdi,%rsi,1), %r9b
    pushq   %r9

    movq    $24, %rdi
    call    malloc
    movq    %rax, %r8
    movq    $3,    (%r8)
    movq    $1,    8(%r8)

    popq    %r9
    movl    $0, %r10d
    shl     $8, %r10
    or      %r10, %r9
    movq    %r9,   16(%r8)

    movq    %r8,  %rdi        
    jmp      P_done

  PERR_index:
    movq    $S_error_index, %rdi
    call    P_runtime_error    

  PERR_type:
    movq    $S_error_type,  %rdi
    call    P_runtime_error    

  P_done:
    leave
    ret


 F_len:
      pushq   %rbp
      movq    %rsp, %rbp
      andq    $-16, %rsp
      movq    (%rdi), %rax    # rax = tag
      cmpq    $3, %rax         # string?
      je      LFL_str
      cmpq    $4, %rax         # list?
      je      LFL_list
      movq    $S_error_type, %rdi
      call    P_runtime_error

  LFL_str:
      movq    8(%rdi), %rax    # string length
      jmp     LFL_done

  LFL_list:
      movq    8(%rdi), %rax    # list length

  LFL_done:
      movq    %rax, %rdi       # prepare to box
      call    P_alloc_int

      movq    %rbp, %rsp
      popq    %rbp
      ret

type_error:
    movq    $S_error_type, %rdi
    call    P_runtime_error

error:
    movq    $S_error_index, %rdi
    call    P_runtime_error

done:
    popq %rbp
    ret
"

let data_inline = "
S_message_int:
  .string    \"%d\"
S_message_string:
  .string    \"%s\"
S_message_None:
  .string    \"None\"
S_message_True:
  .string    \"True\"
S_message_False:
  .string    \"False\"
S_newline:
  .string    \"\\n\"
C_None:
  .quad       0
C_False:
  .quad       1
  .quad       0
C_True:
  .quad       1
  .quad       1
S_lb:
  .string    \"[\"
S_rb:
  .string    \"]\"
S_comma_space:
  .string    \", \"
S_error_fmt:
    .string \"Runtime error: %s\"
S_error_div_zero:
    .string \"division by zero\"
S_error_type:
    .string \"type error\"
S_error_compare_type:
   .string   \"TypeError: cannot compare values of different types\"
S_error_arity:
     .string \"arity mismatch\"
S_error_index:
     .string \"index out of range\"
"

let leave =
  movq (reg rbp) (reg rsp) ++ popq rbp ++ ret

let enter =
  pushq (reg rbp) ++ movq (reg rsp) (reg rbp)

let new_label =
  let c = ref 0 in
  fun () -> incr c; "L_" ^ (string_of_int !c)

type env_alloc = {
  mutable nb_local: int;
  mutable nb_total: int;
}

let empty_alloc_env () = {
  nb_local = 0;
  nb_total = 0;
}

let alloc_var env v =
  if v.v_ofs = -1 then begin
    env.nb_local <- env.nb_local + 1;
    let ofs = -8 * env.nb_local in
    v.v_ofs <- ofs;
    env.nb_total <- env.nb_total + 1
  end

let rec alloc_vars env (s: Ast.tstmt) =
  match s with
  | TSif (_, s1, s2) -> alloc_vars env s1; alloc_vars env s2
  | TSreturn _ -> ()
  | TSassign (x, _) -> alloc_var env x
  | TSprint _ -> ()
  | TSblock sl ->
      List.iter (alloc_vars env) sl
  | TSfor (x, _, s) ->
      alloc_var env x; alloc_vars env s
  | TSwhile(_ , s)  ->
        alloc_vars env s
  | TSeval _ -> ()
  | TSset (_, _, _) -> ()

type string_env = (string, string) Hashtbl.t

let string_env : string_env = Hashtbl.create 16

let new_string =
  let c = ref 0 in
  fun s -> incr c;
    let l = "S_" ^ (string_of_int !c) in
    Hashtbl.add string_env l s;
    l

let get_function_label (f: fn) =
  if f.fn_anoymous
  then "AF_" ^ f.fn_name
  else "F_" ^ f.fn_name

let rec fold_i f i acc l =
  match l with
  | [] -> acc
  | x :: r -> fold_i f (i + 1) (f i acc x) r

let rec compile_expr (e: Ast.texpr) =
  match e with
  | TEcst (Cint n) ->
      movq (imm64 n) (reg rdi) ++ call "P_alloc_int" ++
      movq (reg rax) (reg rdi)
  | TEcst Cnone ->
      movq (ilab "C_None") (reg rdi)
  | TEcst (Cstring s) ->
      let l = new_string s in
      movq (ilab l) (reg rdi)
  | TEcst (Cbool b) ->
      if b then movq (ilab "C_True") (reg rdi)
      else movq (ilab "C_False") (reg rdi)
  | TEvar x ->
      let ofs = x.v_ofs in
      movq (ind ~ofs rbp) (reg rdi)
  | TEfn f ->
      movq (ilab (get_function_label f)) (reg rdi)
  | TEbinop ((Badd|Bsub|Bmul|Bdiv|Bmod as op), e1, e2) ->
        compile_expr e1 ++
        pushq (reg rdi) ++
        compile_expr e2 ++
        movq  (reg rdi) (reg rsi) ++
        popq  rdi               ++
        (match op with
         | Badd -> call "P_add"
         | Bsub -> call "P_sub_int"
         | Bdiv -> call "P_div_int"
         | Bmod -> call "P_mod_int"
         | Bmul ->
             movq (ind ~ofs:8 rdi) (reg rdi) ++
             movq (ind ~ofs:8 rsi) (reg rsi) ++
             imulq (reg rsi)        (reg rdi) ++
             call "P_alloc_int"
         | _ -> assert false) ++
        movq  (reg rax) (reg rdi)    

  | TEbinop ((Beq | Bneq | Blt | Ble | Bgt | Bge as op), te1, te2) ->
      let l_true  = new_label () in
      let l_end   = new_label () in
      compile_expr te1 ++ pushq (reg rdi) ++
      compile_expr te2 ++ popq  rsi ++
      movq  (imm (if op = Beq then 0 else 1)) (reg r10) ++
      call "compare_values" ++
      cmpq (imm 0) (reg rax) ++
      begin match op with
        | Beq  -> je l_true
        | Bneq -> jne l_true
        | Blt -> jg  l_true
        | Ble -> jge l_true
        | Bgt -> jl  l_true
        | Bge -> jle l_true
        | _    -> assert false
      end ++
      movq  (ilab "C_False") (reg rdi) ++
      jmp   l_end ++
      label l_true ++
      movq  (ilab "C_True")  (reg rdi) ++
      label l_end
  | TEbinop (Band, e1, e2) ->
      let l_right = new_label () in
      let l_end   = new_label () in
      compile_expr e1 ++
      movq (reg rdi) (reg r10) ++
      call "P_test" ++
      testq (reg rax) (reg rax) ++
      jnz  l_right ++
      movq (reg r10) (reg rdi) ++  
      jmp  l_end ++
      label l_right ++
      compile_expr e2 ++  
      label l_end
  | TEbinop (Bor, e1, e2) ->
      let l_false = new_label () in
      let l_end   = new_label () in
      compile_expr e1 ++
      movq (reg rdi) (reg r10) ++   
      call "P_test" ++
      testq (reg rax) (reg rax) ++
      jz   l_false ++              
      movq (reg r10) (reg rdi) ++  
      jmp  l_end ++
      label l_false ++
      compile_expr e2 ++
      label l_end
  | TEunop (Uneg, e) ->
      compile_expr (TEcst (Cint 0L)) ++
      pushq (reg rdi) ++
      compile_expr e ++
      movq  (reg rdi) (reg rsi) ++
      popq rdi ++
      call  "P_sub_int" ++
      movq  (reg rax) (reg rdi)
  | TEunop (Unot, e) ->
      let l_true = new_label () in
      let l_end  = new_label () in
      compile_expr e ++
      call    "P_test" ++
      testq   (reg rax) (reg rax) ++
      jz      l_true ++
      movq  (ilab "C_False") (reg rdi) ++
      jmp   l_end     ++
      label   l_true   ++
      movq  (ilab "C_True")  (reg rdi) ++
      label   l_end
  | TEcall (TEfn f, [arg]) when f.fn_name = "len" ->
      let lbl_str   = new_label () in
      let lbl_list  = new_label () in
      let lbl_after = new_label () in
      compile_expr arg ++
      movq (ind ~ofs:0 rdi) (reg rax) ++
      cmpq (imm 3) (reg rax) ++
      je   lbl_str ++
      cmpq (imm 4) (reg rax) ++
      je   lbl_list ++
      movq (ilab "S_error_type") (reg rdi) ++
      call  "P_runtime_error" ++
      label lbl_str ++
      movq (ind ~ofs:8 rdi) (reg rdi) ++
      jmp  lbl_after ++
      label lbl_list ++
      movq (ind ~ofs:8 rdi) (reg rdi) ++
      label lbl_after ++
      call "P_alloc_int" ++        
      movq (reg rax) (reg rdi)   
  | TEcall (fe, el) ->
      let n_args = List.length el in
      let push_arg e =
        compile_expr e ++
        pushq (reg rdi)
      in
      let code_args = List.fold_left (fun acc e -> push_arg e ++ acc) nop el in
      let code_fn = compile_expr fe in
      let expected_args =
        match fe with
        | TEfn f -> List.length f.fn_params
        | _      -> n_args
      in
      let code_arity_check =
        if expected_args <> n_args then
          movq (ilab "S_error_arity") (reg rdi) ++
          call "P_runtime_error"
        else
          nop
      in
      code_args ++
      code_fn ++
      code_arity_check ++
      call_star (reg rdi) ++
      addq (imm (8 * n_args)) (reg rsp) ++
      movq (reg rax) (reg rdi)
  | TElist el ->
      let push_e acc e =
        compile_expr e ++ pushq (reg rdi) ++ acc in
      let pop_e i a e =
        let ofs = 16 + i * 8 in
        a ++ popq rdi ++ movq (reg rdi) (ind ~ofs rax) in
      List.fold_left push_e nop el ++
      movq (imm (List.length el)) (reg rdi) ++
      call "P_alloc_list" ++
      fold_i pop_e 0 nop el ++
      movq (reg rax) (reg rdi)
  
  | TErange e ->
      let lbl_loop  = new_label () in
      let lbl_done  = new_label () in
      let lbl_type  = new_label () in
      let lbl_final = new_label () in
      compile_expr e ++
      movq (ind ~ofs:0 rdi) (reg rax) ++    
      cmpq (imm 2)          (reg rax) ++
      jne   lbl_type ++
      movq (ind ~ofs:8 rdi) (reg rsi) ++   
      movq (reg rsi)        (reg rdi) ++
      call "P_alloc_list" ++
      movq (reg rax)        (reg rbx) ++   
      movq (ind ~ofs:8 rbx) (reg rcx) ++
      leaq (ind ~ofs:16 rbx) r8 ++
      movq (imm 0)          (reg r9) ++
      label lbl_loop ++
      cmpq (reg rcx)        (reg r9) ++
      jge   lbl_done ++
      pushq (reg rbx) ++
      pushq (reg rcx) ++
      pushq (reg r8)  ++
      pushq (reg r9)  ++
      movq (reg r9)    (reg rdi) ++
      call "P_alloc_int"           ++ 
      popq  r9        ++
      popq  r8        ++
      popq  rcx       ++
      popq  rbx       ++
      movq (reg rax) (ind ~ofs:0 r8) ++
      addq (imm 8)   (reg r8)         ++
      incq (reg r9)                  ++
      jmp lbl_loop ++
      label lbl_done ++
      movq (reg rbx) (reg rdi) ++
      jmp lbl_final ++
      label lbl_type ++
      movq (ilab "S_error_type") (reg rdi) ++
      call "P_runtime_error" ++
      label lbl_final
  | TEget (e1, e2) ->
      compile_expr e1           ++  
      pushq    (reg rdi)        ++  
      compile_expr e2           ++ 
      movq     (ind ~ofs:8 rdi) (reg rsi) ++
      popq     ( rdi)        ++ 
      call     "P_index"        ++  
      movq     (reg rdi)       (reg rdi) 
    
let rec compile_stmt exit_lbl (s: Ast.tstmt) =
  match s with
  | TSif (e, s1, s2) ->
      let l_else = new_label () in
      let l_end = new_label () in
      compile_expr e ++
      call "P_test" ++
      testq (reg rax) (reg rax) ++
      jz l_else ++
      compile_stmt exit_lbl s1 ++
      jmp l_end ++
      label l_else ++
      compile_stmt exit_lbl s2 ++
      label l_end
  | TSreturn e ->
      compile_expr e ++
      movq (reg rdi) (reg rax) ++
      jmp exit_lbl
  | TSassign (x, e) ->
      let ofs = x.v_ofs in
      compile_expr e ++ 
      movq (reg rdi) (ind ~ofs rbp)
  | TSprint e ->
      compile_expr e ++ call "P_print" ++ call "P_print_newline"
  | TSblock sl ->
      List.fold_left (fun c s -> c ++ compile_stmt exit_lbl s) nop sl
  | TSfor (v, list_te, body) ->
      let lbl_iter_ok = new_label () in
      let lbl_start   = new_label () in
      let lbl_end     = new_label () in
      compile_expr list_te ++
      movq   (reg rdi)        (reg rbx)    ++
      movq   (ind ~ofs:0 rbx) (reg rax)    ++
      cmpq   (imm 4)          (reg rax)    ++
      je     lbl_iter_ok     ++
      cmpq   (imm 3)          (reg rax)    ++
      je     lbl_iter_ok     ++
      movq   (ilab "S_error_type") (reg rdi) ++
      call   "P_runtime_error"        ++
      label lbl_iter_ok       ++
      movq   (ind ~ofs:8 rbx) (reg r12)    ++
      movq   (imm 0)          (reg rax)    ++
      label lbl_start        ++
      cmpq   (reg r12)        (reg rax)    ++
      jge     lbl_end        ++
      pushq  (reg rax)        ++
      movq   (reg rax)        (reg rsi)    ++
      shlq   (imm 3)          (reg rsi)    ++
      addq   (imm 16)         (reg rsi)    ++
      movq   (reg rbx)        (reg rdi)    ++
      addq   (reg rsi)        (reg rdi)    ++
      movq   (ind ~ofs:0 rdi) (reg rdi)    ++
      movq   (reg rdi)        (ind ~ofs:v.v_ofs rbp) ++
      compile_stmt exit_lbl body  ++
      popq   rax              ++
      incq   (reg rax)        ++
      jmp    lbl_start        ++
      label lbl_end  
  | TSwhile (cond, body) ->
      let l_start = new_label () in
      let l_end   = new_label () in
      label l_start ++
      compile_expr cond ++
      call "P_test" ++
      testq (reg rax) (reg rax) ++
      jz    l_end     ++      
      compile_stmt exit_lbl body ++
      jmp   l_start ++
      label l_end
  | TSeval e ->
      compile_expr e
  | TSset (e1, e2, e3) ->
      compile_expr e1 ++
      movq  (reg rdi)        (reg rbx) ++
      compile_expr e2 ++
      movq  (ind ~ofs:8 rdi) (reg rsi) ++
      movq  (reg rbx)        (reg r12) ++  
      shlq  (imm 3)          (reg rsi) ++  
      addq  (imm 16)         (reg rsi) ++ 
      addq  (reg rsi)        (reg r12) ++  
      compile_expr e3 ++
      movq  (reg rdi)        (ind ~ofs:0 r12)

let compile_tdef (fn, tstmt) =
  let ofs = ref 8 in
  let alloc_param v =
    ofs := !ofs + 8;
    v.v_ofs <- !ofs in
  List.iter alloc_param fn.fn_params;
  let env = empty_alloc_env () in
  alloc_vars env tstmt;
  let exit_lbl = new_label () in
  let body = compile_stmt exit_lbl tstmt in
  let lbl = label (get_function_label fn) in
  let locals =
    if env.nb_total = 0 then nop
    else subq (imm (8 * env.nb_total)) (reg rsp) in
  lbl ++ 
  enter ++ 
  locals ++ 
  body ++ 
  label exit_lbl ++ 
  leave 

let compile_main (fn, tstmt) =
  let env = empty_alloc_env () in
  alloc_vars env tstmt;
  let l = new_label () in
  let locals =
    if env.nb_total = 0 then nop
    else
      subq (imm (8 * env.nb_total)) (reg rsp) in
  enter ++
  locals ++
  compile_stmt l tstmt ++
  xorq (reg rax) (reg rax)

let inline_strings () =
  let alloc_string l s acc =
    label l ++ dquad [3; String.length s] ++ string s ++ acc in
  Hashtbl.fold alloc_string string_env nop

let file ?debug:(b=false) (p: Ast.tfile) : X86_64.program =
  debug := b;
  match p with
  | [] -> assert false
  | m :: r ->
    let cmain = compile_main m ++ leave in
    let cfile =
      List.fold_left (fun a td -> a ++ compile_tdef td) nop r in
    { text = globl "main" ++ label "main" ++ cmain ++
             cfile ++ inline preamble;
      data = inline data_inline ++ inline_strings (); }