open Syntax

let rec insert_checks (e : 'a Ast.t) : 'a Ast.t =
  let lift desc = { e with desc } in
  (match e.desc with
  | Const const -> Const const
  | Var x -> Var x
  | Lam (x, e) -> Lam (x, insert_checks e)
  | App (e1, exprs) -> App (insert_checks e1, List.map insert_checks exprs)
  | Prim1 (op1, e) ->
      let ty, _ = Type.ty_of_op1 op1 in
      let e' = lift (Ast.Ascribe (insert_checks e, ty)) in
      Prim1 (op1, e')
  | Prim2 (op2, e1, e2) ->
      let ty1, ty2, _ = Type.ty_of_op2 op2 in
      let e1' = lift (Ast.Ascribe (insert_checks e1, ty1)) in
      let e2' = lift (Ast.Ascribe (insert_checks e2, ty2)) in
      Prim2 (op2, e1', e2')
  | If (e1, e2, e3) ->
      let e1' = lift (Ast.Ascribe (insert_checks e1, Type.Bool)) in
      let e2' = insert_checks e2 in
      let e3' = insert_checks e3 in
      If (e1', e2', e3')
  | Let (binding, e) ->
      let binding' =
        Ast.{ name = binding.name; term = insert_checks binding.term }
      in
      let e' = insert_checks e in
      Let (binding', e')
  | Letrec (bindings, e) ->
      let bindings' =
        List.map
          (fun binding ->
            Ast.{ binding with term = insert_checks binding.term })
          bindings
      in
      let e' = insert_checks e in
      Letrec (bindings', e')
  | Tuple exprs -> Tuple (List.map insert_checks exprs)
  | Get (e, i) -> Get (insert_checks e, i)
  | Vect (e1, e2) ->
      let e1' = lift (Ast.Ascribe (insert_checks e1, Type.Int)) in
      Vect (e1', e2)
  | VectGet (e1, e2) ->
      let e1' = lift (Ast.Ascribe (insert_checks e1, Type.Vec Type.Any)) in
      let e2' = lift (Ast.Ascribe (insert_checks e2, Type.Int)) in
      Vect (e1', e2')
  | VectSet (e1, e2, e3) ->
      let e1' = lift (Ast.Ascribe (insert_checks e1, Type.Vec Type.Any)) in
      let e2' = lift (Ast.Ascribe (insert_checks e2, Type.Int)) in
      let e3' = insert_checks e3 in
      VectSet (e1', e2', e3')
  | VectLength e ->
      let e' = lift (Ast.Ascribe (insert_checks e, Type.Vec Type.Any)) in
      VectLength e'
  | Ascribe (e, ty) -> Ascribe (insert_checks e, ty)
  | For (x, e1, e2, e3) ->
      let e1' = lift (Ast.Ascribe (insert_checks e1, Type.Int)) in
      let e2' = lift (Ast.Ascribe (insert_checks e2, Type.Int)) in
      let e3' = insert_checks e3 in
      For (x, e1', e2', e3')
  | Variant (x, exprs, ty) ->
      let tys =
        match ty with
        | Type.Rec (_, Type.Variant cases) | Type.Variant cases ->
            let case =
              List.find (fun (Type.Case (name, _)) -> name = x) cases
            in
            let (Type.Case (_, tys)) = case in
            tys
        | _ -> failwith "Expected variant type"
      in
      let exprs' =
        List.map2
          (fun e ty -> lift (Ast.Ascribe (insert_checks e, ty)))
          exprs tys
      in
      Variant (x, exprs', ty)
  | Match (e, patterns, ty) ->
      let e' = lift (Ast.Ascribe (insert_checks e, ty)) in
      let patterns' =
        patterns
        |> List.map (fun pat -> Ast.{ pat with body = insert_checks pat.body })
      in
      Match (e', patterns', ty)
  | Seq exprs -> Seq (List.map insert_checks exprs))
  |> lift
