-- Juvix/Core/Main/Semantics/Eval/Indexed.lean

import Juvix.Core.Main.Semantics.Eval

import Mathlib.Tactic.Linarith

namespace Juvix.Core.Main

inductive Eval.Indexed : Nat → Env → Expr → Value → Prop where
  | var {n env name idx val} :
    env[idx]? = some (Object.value val) →
    Eval.Indexed n env (Expr.var name idx) val
  | var_rec {n env name idx env' expr val} :
    env[idx]? = some (Object.delayed env' expr) →
    Eval.Indexed n env' expr val →
    Eval.Indexed n env (Expr.var name idx) val
  | unit {n env} :
    Eval.Indexed n env Expr.unit Value.unit
  | const {n env c} :
    Eval.Indexed n env (Expr.const c) (Value.const c)
  | constr {n env name} :
    Eval.Indexed n env (Expr.constr name) (Value.constr_app name [])
  | app {n n₁ n₂ env env' f body arg val val'} :
    n₁ + n₂ + 1 ≤ n →
    Eval.Indexed n₁ env f (Value.closure env' body) →
    Eval.Indexed (n₁ + 1) env arg val →
    Eval.Indexed n₂ (val ∷ env') body val' →
    Eval.Indexed n env (Expr.app f arg) val'
  | constr_app {n n' env ctr ctr_name ctr_args_rev arg val} :
    n' < n →
    Eval.Indexed n env ctr (Value.constr_app ctr_name ctr_args_rev) →
    Eval.Indexed n' env arg val →
    Eval.Indexed n env (Expr.constr_app ctr arg) (Value.constr_app ctr_name (val :: ctr_args_rev))
  | binop {n env op arg₁ arg₂ val₁ val₂} :
    Eval.Indexed n env arg₁ (Value.const (Constant.int val₁)) →
    Eval.Indexed n env arg₂ (Value.const (Constant.int val₂)) →
    Eval.Indexed n env (Expr.binop op arg₁ arg₂) (Value.const (Constant.int (eval_binop_int op val₁ val₂)))
  | lambda {n env name body} :
    Eval.Indexed n env (Expr.lambda name body) (Value.closure env body)
  | save {n n₁ n₂ env name value body val val'} :
    n₁ + n₂ ≤ n →
    Eval.Indexed n₁ env value val →
    Eval.Indexed n₂ (val ∷ env) body val' →
    Eval.Indexed n env (Expr.save name value body) val'
  | branch_matches {n n' env name args_rev body val} :
    n' < n →
    Eval.Indexed n' (args_rev.map Object.value ++ env) body val →
    Eval.Indexed n (Value.constr_app name args_rev ∷ env) (Expr.branch name _ body _) val
  | branch_fails {n env name name' args_rev next val} :
    name ≠ name' →
    Eval.Indexed n (Value.constr_app name args_rev ∷ env) next val →
    Eval.Indexed n (Value.constr_app name args_rev ∷ env) (Expr.branch name' _ _ next) val
  | recur {n n' env name body val} :
    n' < n →
    Eval.Indexed n' (Object.delayed env (Expr.recur name body) :: env) body val →
    Eval.Indexed n env (Expr.recur name body) val

notation:40 env:40 " ⊢ " e:40 " ↦(" n ") " v:40 => Eval.Indexed n env e v

-- Repository definitions used:
content: syntax:100 expr:100 ppSpace expr:101 : expr

syntax:50 expr:50 " + " expr:51 : expr

syntax:50 expr:50 " - " expr:50 : expr

syntax:60 expr:60 " * " expr:61 : expr

syntax:60 expr:60 " / " expr:60 : expr

syntax "case " expr " of " cases " end" : expr

content: notation:40 env:40 " ⊢ " e:40 " ↦ " v:40 => Eval env e v

content: notation:40 env:40 " ⊢ " e:40 " ↓" => Eval.Defined env e

content: notation:40 env:40 " ⊢ " e:40 " ↦(" n ") " v:40 => Eval.Indexed n env e v

content: macro_rules
  | `(⟪$s:ident ♯ $i:num⟫) => `(Expr.var $(Lean.Syntax.mkStrLit s.getId.toString) $i)
  | `(⟪$num:num⟫) => `(Expr.const (Constant.int $num))
  | `(⟪$s:str⟫) => `(Expr.const (Constant.string $s))
  | `(⟪υ⟫) => `(Expr.unit)
  | `(⟪$e₁:expr $e₂:expr⟫) => `(Expr.app ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪λ $s:ident $e:expr⟫) => `(Expr.lambda $(Lean.Syntax.mkStrLit s.getId.toString) ⟪$e⟫)
  | `(⟪λ $ss:ident* . $e:expr⟫) => mkLambdas ss e
  | `(⟪ $s:ident ⟫) => `(Expr.constr $(Lean.Syntax.mkStrLit s.getId.toString))
  | `(⟪ $s:ident $es:expr* ⟫) => mkConstrApp s es
  | `(⟪$e₁ + $e₂⟫) => `(Expr.binop BinaryOp.add_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪$e₁ - $e₂⟫) => `(Expr.binop BinaryOp.sub_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪$e₁ * $e₂⟫) => `(Expr.binop BinaryOp.mul_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪$e₁ / $e₂⟫) => `(Expr.binop BinaryOp.div_int ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪let $s:ident := $e₁:expr in $e₂:expr⟫) => `(Expr.save $(Lean.Syntax.mkStrLit s.getId.toString) ⟪$e₁⟫ ⟪$e₂⟫)
  | `(⟪letrec $s:ident := $e₁:expr in $e₂:expr⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.save $name (Expr.recur $name ⟪$e₁⟫) ⟪$e₂⟫)
  | `(⟪rec $s:ident $e:expr⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.recur $name ⟪$e⟫)
  | `(⟪⊥⟫) => `(Expr.fail)
  | `(⟪cases| | $s:ident => $e:expr $cs:cases ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.branch $name [] ⟪$e⟫ ⟪cases|$cs⟫)
  | `(⟪cases| | $s:ident $ss:ident* => $e:expr $cs:cases ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    let vnames := ss.toList.map fun s => s.getId.toString
    `(Expr.branch $name $(quote vnames) ⟪$e⟫ ⟪cases|$cs⟫)
  | `(⟪cases| | $s:ident => $e:expr ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    `(Expr.branch $name [] ⟪$e⟫ Expr.fail)
  | `(⟪cases| | $s:ident $ss:ident* => $e:expr ⟫) =>
    let name := Lean.Syntax.mkStrLit s.getId.toString
    let vnames := ss.toList.map fun s => s.getId.toString
    `(Expr.branch $name $(quote vnames) ⟪$e⟫ Expr.fail)
  | `(⟪cases| | _ => $e:expr ⟫) =>
    `(⟪$e⟫)
  | `(⟪case $e:expr of $cs:cases end⟫) => do
    `(Expr.save "_case_" ⟪$e⟫ ⟪cases|$cs⟫)
  | `(⟪($e)⟫) => `(⟪$e⟫)

content: @[aesop unsafe constructors]
inductive Eval : Env → Expr → Value → Prop where
  | var {env name idx val} :
    env[idx]? = some (Object.value val) →
    Eval env (Expr.var name idx) val
  | var_rec {env name idx env' expr val} :
    env[idx]? = some (Object.delayed env' expr) →
    Eval env' expr val →
    Eval env (Expr.var name idx) val
  | unit {env} :
    Eval env Expr.unit Value.unit
  | const {env c} :
    Eval env (Expr.const c) (Value.const c)
  | constr {env name} :
    Eval env (Expr.constr name) (Value.constr_app name [])
  | app {env env' f body arg val val'} :
    Eval env f (Value.closure env' body) →
    Eval env arg val →
    Eval (val ∷ env') body val' →
    Eval env (Expr.app f arg) val'
  | constr_app {env ctr ctr_name ctr_args_rev arg val} :
    Eval env ctr (Value.constr_app ctr_name ctr_args_rev) →
    Eval env arg val →
    Eval env (Expr.constr_app ctr arg) (Value.constr_app ctr_name (val :: ctr_args_rev))
  | binop {env op arg₁ arg₂ val₁ val₂} :
    Eval env arg₁ (Value.const (Constant.int val₁)) →
    Eval env arg₂ (Value.const (Constant.int val₂)) →
    Eval env (Expr.binop op arg₁ arg₂) (Value.const (Constant.int (eval_binop_int op val₁ val₂)))
  | lambda {env name body} :
    Eval env (Expr.lambda name body) (Value.closure env body)
  | save {env name value body val val'} :
    Eval env value val →
    Eval (val ∷ env) body val' →
    Eval env (Expr.save name value body) val'
  | branch_matches {env name args_rev body val} :
    Eval (args_rev.map Object.value ++ env) body val →
    Eval (Value.constr_app name args_rev ∷ env) (Expr.branch name _ body _) val
  | branch_fails {env name name' args_rev next val} :
    name ≠ name' →
    Eval (Value.constr_app name args_rev ∷ env) next val →
    Eval (Value.constr_app name args_rev ∷ env) (Expr.branch name' _ _ next) val
  | recur {env name body val} :
    Eval (Object.delayed env (Expr.recur name body) :: env) body val →
    Eval env (Expr.recur name body) val

content: inductive BinaryOp : Type where
  | add_int : BinaryOp
  | sub_int : BinaryOp
  | mul_int : BinaryOp
  | div_int : BinaryOp
  deriving Inhabited, BEq, DecidableEq

content: def eval_binop_int (op : BinaryOp) (i₁ i₂ : Int) : Int :=
  match op with
  | BinaryOp.add_int => i₁ + i₂
  | BinaryOp.sub_int => i₁ - i₂
  | BinaryOp.mul_int => i₁ * i₂
  | BinaryOp.div_int => i₁ / i₂

content: inductive Expr : Type where
  | var : (name : String) → (index : Nat) → Expr
  | unit : Expr
  | const : Constant → Expr
  | constr : Name → Expr
  | app : Expr → Expr → Expr
  | constr_app : Expr → Expr → Expr
  | binop : (oper : BinaryOp) → (arg₁ arg₂ : Expr) → Expr
  | lambda : (var_name : String) → (body : Expr) → Expr
  | save : (var_name : String) → (value : Expr) → (body : Expr) → Expr
  | branch : (constr : Name) → (var_names : List Name) → (body : Expr) → (next : Expr) → Expr
  | recur : (var_name : Name) → (body : Expr) → Expr
  | fail : Expr
  deriving Inhabited, BEq, DecidableEq

content: abbrev cons_value (v : Value) (env : Env) : Env := Object.value v :: env

  inductive Object : Type where
    | value : Value → Object
    | delayed : (env : List Object) → Expr → Object
    deriving Inhabited

  inductive Value : Type where
    | unit : Value
    | const : Constant → Value
    | constr_app : (constr : Name) → (args_rev : List Value) → Value
    | closure : (env : List Object) → (value : Expr) → Value
    deriving Inhabited

content: inductive Constant : Type where
  | int : Int → Constant
  | string : String → Constant
  deriving Inhabited, BEq, DecidableEq

content: abbrev Name : Type := String

content: abbrev Env : Type := List Object

content: inductive Context : Type where
  | hole : Context
  | app_left : Context → Expr → Context
  | app_right : Expr → Context → Context
  | constr_app_left : Context → Expr → Context
  | constr_app_right : Expr → Context → Context
  | binop_left : (oper : BinaryOp) → (arg₁ : Context) → (arg₂ : Expr) → Context
  | binop_right : (oper : BinaryOp) → (arg₁ : Expr) → (arg₂ : Context) → Context
  | lambda : (var_name : String) → (body : Context) → Context
  | save_left : (var_name : String) → (value : Context) → (body : Expr) → Context
  | save_right : (var_name : String) → (value : Expr) → (body : Context) → Context
  | branch_left : (constr : Name) → (var_names : List Name) → (body : Context) → (next : Expr) → Context
  | branch_right : (constr : Name) → (var_names : List Name) → (body : Expr) → (next : Context) → Context
  | recur : (var_name : Name) → (ctx : Context) → Context
  deriving Inhabited, BEq

content: infixr:50 " ∷ " => cons_value

content: notation:40 env:40 " ⊢ " e:40 " ↦ " v:40 => Eval env e v

content: infixl:100 " @@ " => Expr.app

-- Repository lemmas used:
content: lemma Termination.var {env name i v} :
  env[i]? = some (Object.value v) →
  env ⊢ Expr.var name i ↓

content: lemma Termination.var_rec {env name i env' e} :
    env[i]? = some (Object.delayed env' e) →
    env' ⊢ e ↓ →
    env ⊢ (Expr.var name i) ↓

content: lemma Termination.unit {env} :
  env ⊢ Expr.unit ↓

content: lemma Termination.const {env c} :
  env ⊢ Expr.const c ↓

content: lemma Termination.constr {env nm} :
  env ⊢ Expr.constr nm ↓

content: lemma Termination.app {env e₁ e₂ v env' body} :
  env ⊢ e₁ ↦ Value.closure env' body →
  env ⊢ e₂ ↦ v →
  v ∷ env' ⊢ body ↓ →
  env ⊢ Expr.app e₁ e₂ ↓

content: lemma Termination.constr_app {env e₁ e₂ e₁₁ e₁₂} :
  env ⊢ e₁ ↦ Value.constr_app e₁₁ e₁₂ →
  env ⊢ e₂ ↓ →
  env ⊢ Expr.constr_app e₁ e₂ ↓

content: lemma Termination.binop {env op e₁ e₂ c₁ c₂} :
  env ⊢ e₁ ↦ Value.const (Constant.int c₁) →
  env ⊢ e₂ ↦ Value.const (Constant.int c₂) →
  env ⊢ Expr.binop op e₁ e₂ ↓

content: lemma Termination.lambda {env name e} : env ⊢ Expr.lambda name e ↓

content: lemma Termination.save {env name e₁ e₂ v} :
  env ⊢ e₁ ↦ v →
  v ∷ env ⊢ e₂ ↓ →
  env ⊢ Expr.save name e₁ e₂ ↓

content: lemma Termination.branch_matches {env name vnames args_rev e e'} :
  args_rev.map Object.value ++ env ⊢ e ↓ →
  Value.constr_app name args_rev ∷ env ⊢ Expr.branch name vnames e e' ↓

content: lemma Termination.branch_fails {env name name' vnames args_rev e e'} :
  name ≠ name' →
  Value.constr_app name args_rev ∷ env ⊢ e' ↓ →
  Value.constr_app name args_rev ∷ env ⊢ Expr.branch name' vnames e e' ↓

content: lemma Termination.recur {env name e} :
  Object.delayed env (Expr.recur name e) :: env ⊢ e ↓ →
  env ⊢ Expr.recur name e ↓

-- Local definitions used:
content: inductive Eval.Indexed : Nat → Env → Expr → Value → Prop where
  | var {n env name idx val} :
    env[idx]? = some (Object.value val) →
    Eval.Indexed n env (Expr.var name idx) val
  | var_rec {n env name idx env' expr val} :
    env[idx]? = some (Object.delayed env' expr) →
    Eval.Indexed n env' expr val →
    Eval.Indexed n env (Expr.var name idx) val
  | unit {n env} :
    Eval.Indexed n env Expr.unit Value.unit
  | const {n env c} :
    Eval.Indexed n env (Expr.const c) (Value.const c)
  | constr {n env name} :
    Eval.Indexed n env (Expr.constr name) (Value.constr_app name [])
  | app {n n₁ n₂ env env' f body arg val val'} :
    n₁ + n₂ + 1 ≤ n →
    Eval.Indexed n₁ env f (Value.closure env' body) →
    Eval.Indexed (n₁ + 1) env arg val →
    Eval.Indexed n₂ (val ∷ env') body val' →
    Eval.Indexed n env (Expr.app f arg) val'
  | constr_app {n n' env ctr ctr_name ctr_args_rev arg val} :
    n' < n →
    Eval.Indexed n env ctr (Value.constr_app ctr_name ctr_args_rev) →
    Eval.Indexed n' env arg val →
    Eval.Indexed n env (Expr.constr_app ctr arg) (Value.constr_app ctr_name (val :: ctr_args_rev))
  | binop {n env op arg₁ arg₂ val₁ val₂} :
    Eval.Indexed n env arg₁ (Value.const (Constant.int val₁)) →
    Eval.Indexed n env arg₂ (Value.const (Constant.int val₂)) →
    Eval.Indexed n env (Expr.binop op arg₁ arg₂) (Value.const (Constant.int (eval_binop_int op val₁ val₂)))
  | lambda {n env name body} :
    Eval.Indexed n env (Expr.lambda name body) (Value.closure env body)
  | save {n n₁ n₂ env name value body val val'} :
    n₁ + n₂ ≤ n →
    Eval.Indexed n₁ env value val →
    Eval.Indexed n₂ (val ∷ env) body val' →
    Eval.Indexed n env (Expr.save name value body) val'
  | branch_matches {n n' env name args_rev body val} :
    n' < n →
    Eval.Indexed n' (args_rev.map Object.value ++ env) body val →
    Eval.Indexed n (Value.constr_app name args_rev ∷ env) (Expr.branch name _ body _) val
  | branch_fails {n env name name' args_rev next val} :
    name ≠ name' →
    Eval.Indexed n (Value.constr_app name args_rev ∷ env) next val →
    Eval.Indexed n (Value.constr_app name args_rev ∷ env) (Expr.branch name' _ _ next) val
  | recur {n n' env name body val} :
    n' < n →
    Eval.Indexed n' (Object.delayed env (Expr.recur name body) :: env) body val →
    Eval.Indexed n env (Expr.recur name body) val

-- Local lemmas used:
content: lemma Eval.Indexed.monotone {n n' env e v} (h : env ⊢ e ↦(n) v) (h' : n ≤ n') : env ⊢ e ↦(n') v

-- Target theorem and proof:
lemma Eval.toIndexed {env e v} (h : env ⊢ e ↦ v) : ∃ n, env ⊢ e ↦(n) v
:= by
  induction h
  case var =>
    exists 0; apply Eval.Indexed.var; assumption
  case var_rec ih =>
    obtain ⟨n, ih⟩ := ih
    exists n
    apply Eval.Indexed.var_rec <;> assumption
  case unit =>
    exists 0; apply Eval.Indexed.unit
  case const =>
    exists 0; apply Eval.Indexed.const
  case constr =>
    exists 0; apply Eval.Indexed.constr
  case app env env' f body arg val val' _ _ _ ih1 ih2 ih3 =>
    obtain ⟨n₁, ih₁⟩ := ih1
    obtain ⟨n₂, ih₂⟩ := ih2
    obtain ⟨n₃, ih₃⟩ := ih3
    have ih₁' : env ⊢ f ↦(n₁ + n₂) Value.closure env' body := by
      apply Eval.Indexed.monotone ih₁; linarith
    have ih₂' : env ⊢ arg ↦(n₁ + n₂ + 1) val := by
      apply Eval.Indexed.monotone ih₂; linarith
    exists (n₁ + n₂ + n₃ + 1)
    apply Eval.Indexed.app
    · rfl
    · exact ih₁'
    · exact ih₂'
    · exact ih₃
  case constr_app ih₁ ih₂ =>
    obtain ⟨n₁, ih₁⟩ := ih₁
    obtain ⟨n₂, ih₂⟩ := ih₂
    exists (n₁ + n₂ + 1)
    apply Eval.Indexed.constr_app (n' := n₁ + n₂)
    · linarith
    · apply Eval.Indexed.monotone ih₁; linarith
    · apply Eval.Indexed.monotone ih₂; linarith
  case binop ih₁ ih₂ =>
    obtain ⟨n₁, ih₁⟩ := ih₁
    obtain ⟨n₂, ih₂⟩ := ih₂
    exists (n₁ + n₂)
    apply Eval.Indexed.binop (n := n₁ + n₂)
    · apply Eval.Indexed.monotone ih₁; linarith
    · apply Eval.Indexed.monotone ih₂; linarith
  case lambda =>
    exists 0; apply Eval.Indexed.lambda
  case save ih₁ ih₂ =>
    obtain ⟨n₁, ih₁⟩ := ih₁
    obtain ⟨n₂, ih₂⟩ := ih₂
    exists (n₁ + n₂)
    apply Eval.Indexed.save (n := n₁ + n₂)
    · rfl
    · apply Eval.Indexed.monotone ih₁; linarith
    · apply Eval.Indexed.monotone ih₂; linarith
  case branch_matches ih =>
    obtain ⟨n, ih⟩ := ih
    exists (n + 1)
    apply Eval.Indexed.branch_matches (n' := n)
    · linarith
    · assumption
  case branch_fails ih =>
    obtain ⟨n, ih⟩ := ih
    exists n
    apply Eval.Indexed.branch_fails <;> assumption
  case recur ih =>
    obtain ⟨n, ih⟩ := ih
    exists n.succ
    apply Eval.Indexed.recur (n' := n)
    · linarith
    · assumption