-- SSA/Projects/Scf/ScfFunctor.lean

import Mathlib.Logic.Function.Iterate

import Mathlib.Tactic.Linarith

import LeanMLIR.Framework

import LeanMLIR.Tactic

import LeanMLIR.ErasedContext

import LeanMLIR.Util

open LeanMLIR

open Ctxt(Var)

namespace ScfFunctor

open TyDenote

class HasTy (d : Dialect) (DenotedTy : Type) [TyDenote d.Ty] [DialectSignature d] where
    ty : d.Ty
    denote_eq : toType ty = DenotedTy := by admit /- proof elided -/

inductive Scf.Op (Op' Ty' : Type) (m') [TyDenote Ty'] [DialectSignature ⟨Op', Ty', m'⟩]
    [DialectDenote ⟨Op', Ty', m'⟩] : Type _
  | coe (o : Op')
  | iterate (k : ℕ) 
  | run (inputty : Ty')  
  | if (inputty retty' : Ty')  
  | for (ty : Ty')
  deriving DecidableEq, Repr

def Scf (d : Dialect) [TyDenote d.Ty] [DialectSignature d] [DialectDenote d] : Dialect where
  Op := Scf.Op d.Op d.Ty d.m
  Ty := d.Ty
  m  := d.m

namespace Scf

section InheritedInstances

variable {d : Dialect} [TyDenote d.Ty] [DialectSignature d] [DialectDenote d]

end InheritedInstances

abbrev LoopBody (t : Type) : Type := Int → t → t

namespace LoopBody

def counterDecorator (δ : Int) (f : LoopBody α) : Int × α → Int × α :=
  fun (i, v) => (i + δ, f i v)

def atZero (f : LoopBody t) : t → t := fun v => f 0 v

end LoopBody

namespace LoopBody.IndexInvariant

end LoopBody.IndexInvariant

namespace LoopBody.counterDecorator

end LoopBody.counterDecorator

variable [TyDenote d.Ty] [DialectSignature d] [DialectDenote d]
  [B : HasBool d] [N : HasNat d] [Z : HasInt d]

open Ctxt (Valuation) in

end Scf

namespace Arith

inductive Ty
| int
| bool
| nat
 deriving DecidableEq, Repr

inductive Op
  | add : Op  
  | add_nat : Op  
  | axpy : Op  
  | neg : Op  
  | const : (val : ℤ) → Op
  | const_nat : (val : ℕ) → Op

abbrev Arith : Dialect := {Op, Ty}

end Arith

abbrev ScfArith := Scf Arith

open LeanMLIR.SingleReturnCompat (Com Expr)

@[simp_denote] def cst (n : ℤ) : Expr ScfArith Γ .pure .int  :=
  Expr.mk
    (op := .coe <| .const n)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

@[simp_denote] def add {Γ : Ctxt _} (e₁ e₂ : Var Γ .int) : Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .add)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

@[simp_denote] def add_nat (e₁ e₂ : Var Γ .nat) : Expr ScfArith Γ .pure .nat :=
  Expr.mk
    (op := .coe <| .add_nat)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

@[simp_denote] def axpy {Γ : Ctxt _} (a : Var Γ .int) (x : Var Γ .nat) (b: Var Γ .int) :
    Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .axpy)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons x <| .cons b .nil)
    (regArgs := .nil)

@[simp_denote] def neg {Γ : Ctxt _} (a : Var Γ .int) : Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .neg)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

@[simp_denote] def iterate {Γ : Ctxt _} (k : Nat) (input : Var Γ Arith.Ty.int)
    (body : Com ScfArith ⟨[.int]⟩ .impure .int) : Expr ScfArith Γ .impure .int :=
  Expr.mk
    (op := .iterate k)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons input .nil)
    (regArgs := HVector.cons body HVector.nil)

@[simp_denote]
def run {Γ : Ctxt _} {t : Arith.Ty} (v : Var Γ t) (body : Com ScfArith ⟨[t]⟩ .impure t) :
    Expr ScfArith Γ .impure t :=
  Expr.mk
    (op := .run t)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons v .nil)
    (regArgs := HVector.cons body <| HVector.nil)

@[simp_denote] def for_ {Γ : Ctxt Arith.Ty} {t : Arith.Ty}
    (start step : Var Γ Arith.Ty.int)
    (niter : Var Γ Arith.Ty.nat) (v : Var Γ t) (body : Com ScfArith ⟨[.int, t]⟩ .impure t) :
      Expr ScfArith Γ .impure t :=
  Expr.mk
    (op := .for t)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons start <| .cons step <| .cons niter <| .cons v .nil)
    (regArgs := HVector.cons body <| HVector.nil)

namespace ForAddToMul

def lhs (vincrement : ℤ) : Com ScfArith ⟨[  .nat,   .int]⟩ .impure .int :=
    Com.letPure (cst 0) <|
    Com.letPure  (cst 1) <|
    Com.var (for_ (t := .int)
                        ⟨  1, rfl⟩
                        ⟨  0, rfl⟩
                        ⟨  2, rfl⟩
                        ⟨  3, rfl⟩ (
      Com.letPure (cst vincrement) <|
      Com.letPure (add ⟨0, rfl⟩ ⟨2, rfl⟩) 
      <| Com.ret ⟨0, rfl⟩)) <|
  Com.ret ⟨0, rfl⟩

def rhs (vincrement : ℤ) : Com ScfArith ⟨[  .nat,   .int]⟩ .pure .int :=
  Com.var (cst vincrement) <|
  Com.var (axpy ⟨0, rfl⟩ ⟨1, rfl⟩ ⟨2, rfl⟩) <|
  Com.ret ⟨0, rfl⟩

open Scf in

open Arith in

-- Repository definitions used:
content: macro "simp_peephole" loc:(location)? : tactic =>
  `(tactic|(
      
      first
      | rw [funext_iff (α := Ctxt.Valuation _)] $[$loc]?
      | change ∀ (_ : Ctxt.Valuation _), _ $[$loc]?
      | skip

      
      simp (config := {failIfUnchanged := false}) only
        [Expr.denote_castPureToEff, simp_denote] $[$loc]?
      
      
      
  ))

content: @[match_pattern]
def cons (hd : Ty) : Ctxt Ty → Ctxt Ty
| ⟨tl⟩ => ⟨hd :: tl⟩

content: structure Ctxt (Ty : Type) : Type where
  ofList :: toList : List Ty
  
  deriving Repr, Lean.ToExpr, DecidableEq

content: def Var (Γ : Ctxt Ty) (t : Ty) : Type :=
  { i : Nat // Γ[i]? = some t }

content: protected instance Expr.decidableEq [DecidableEq d.Op] [DecidableEq d.Ty] :
    {Γ : Ctxt d.Ty} → {ty : List d.Ty} → DecidableEq (Expr d Γ eff ty)
  | Γ, _, .mk op₁ rfl eff_le₁ arg₁ regArgs₁, .mk op₂ eq eff_le₂ arg₂ regArgs₂ =>
    if ho : op₁ = op₂ then by
      subst ho
      letI := HVector.decidableEq arg₁ arg₂
      letI := HVector.decidableEqReg regArgs₁ regArgs₂
      exact decidable_of_iff (arg₁ = arg₂ ∧ regArgs₁ = regArgs₂) (by admit /- proof elided -/
      )
    else isFalse (by admit /- proof elided -/
    )

content: def Var (Γ : Context Ty) (ty : Ty) : Type := { v : VarName // Γ.hasType v ty }

content: protected instance HVector.decidableEqReg [DecidableEq d.Op] [DecidableEq d.Ty] :
    ∀ {l : RegionSignature d.Ty}, DecidableEq (HVector (fun t => Com d t.1 .impure t.2) l)
  | _, .nil, .nil => isTrue rfl
  | _, .cons x₁ v₁, .cons x₂ v₂ =>
    letI := HVector.decidableEqReg v₁ v₂
    letI := Com.decidableEq x₁ x₂
    decidable_of_iff (x₁ = x₂ ∧ v₁ = v₂) (by admit /- proof elided -/
    )

content: inductive Com : Ctxt d.Ty → EffectKind → List d.Ty → Type where
  | rets {Γ} {tys} {eff : EffectKind} (vs : HVector Γ.Var tys) : Com Γ eff tys
  | var (e : Expr Γ eff ty) (body : Com (ty ++ Γ) eff β) : Com Γ eff β

content: structure FlatCom (d : Dialect) [DialectSignature d]  (Γ_in : Ctxt d.Ty) (eff : EffectKind)
    (Γ_out : Ctxt d.Ty) (ts : List d.Ty) where
  lets : Lets d Γ_in eff Γ_out
  rets : HVector Γ_out.Var ts

content: inductive Lets (Γ_in : Ctxt d.Ty) (eff : EffectKind) :
    (Γ_out : Ctxt d.Ty) → Type where
  | nil : Lets Γ_in eff Γ_in
  | var (body : Lets Γ_in eff Γ_out) (e : Expr d Γ_out eff t) : Lets Γ_in eff e.outContext

content: inductive EffectKind
| pure 
| impure 
deriving Repr, DecidableEq, Lean.ToExpr

content: class DialectDenote (d : Dialect) [TyDenote d.Ty] [DialectSignature d] where
  denote : (op : d.Op) → HVector toType (DialectSignature.sig op) →
    (HVector (fun t : Ctxt d.Ty × List d.Ty =>
      t.1.Valuation
      → EffectKind.impure.toMonad d.m (HVector toType t.2))
            (DialectSignature.regSig op)) →
    ((DialectSignature.effectKind op).toMonad d.m
      (HVector toType <| DialectSignature.returnTypes op))

content: def sig          := Signature.sig ∘ s.signature

content: structure Signature (Ty : Type) where
  mkEffectful ::
  sig         : List Ty
  regSig      : RegionSignature Ty
  returnTypes : List Ty
  effectKind  : EffectKind := .pure

content: class DialectSignature (d : Dialect) where
  signature : d.Op → Signature d.Ty

content: def regSig       := Signature.regSig ∘ s.signature

content: abbrev RegionSignature Ty := List (Ctxt Ty × List Ty)

content: inductive Ty (q : Nat) (n : Nat)
  | index : Ty q n
  | integer : Ty q n
  | tensor : Ty q n
  | polynomialLike : Ty q n
  deriving DecidableEq, Repr

content: def effectKind   := Signature.effectKind ∘ s.signature

content: def returnTypes  := Signature.returnTypes ∘ s.signature

content: structure Dialect where
  (Op : Type)
  (Ty : Type)
  (m : Type → Type := Id)

content: inductive Op (q : Nat) (n : Nat)
  | add : Op q n
  | sub : Op q n
  | mul : Op q n
  | mul_constant : Op q n
    
    
  | leading_term : Op q n
  | monomial : Op q n
  | monomial_mul : Op q n
  | from_tensor : Op q n
  | to_tensor : Op q n
  | const (c : R q n) : Op q n
  | const_int (c : Int) : Op q n
  | const_idx (i : Nat) : Op q n

content: class TyDenote (β : Type) : Type 1 where
  toType : β → Type

content: @[simp]
instance Valuation.instAppendHVector (Γ : Ctxt Ty) (ts : List Ty) :
    HAppend (HVector toType ts) (Valuation Γ) (Valuation <| ⟨ts⟩ ++ Γ) where
  hAppend vals V :=
    (Valuation.ofHVector vals) ++ V

content: inductive Ty
  | int
   
  | int2
  deriving DecidableEq, Lean.ToExpr

content: inductive Op
  | noop
  | mkPair
  | unPair
  deriving Lean.ToExpr

content: def Com.ret {Γ : Ctxt d.Ty} {ty : d.Ty} {eff : EffectKind} : Γ.Var ty → Com d Γ eff [ty] :=
  (Com.rets [·]ₕ)

content: def Com.letPure (e : Expr d Γ .pure t) (body : Com d (e.outContext) eff u) : Com d Γ eff u :=
  body.var (e.castPureToEff eff)

content: abbrev Expr.outContext (_ : Expr d Γ eff ts) : Ctxt d.Ty :=
  ts ++ Γ

content: def Expr.castPureToEff (eff : EffectKind) : Expr d Γ .pure t → Expr d Γ eff t :=
  changeEffect (EffectKind.pure_le eff)

content: def Expr.changeEffect {eff₁ eff₂ : EffectKind} (h : eff₁ ≤ eff₂) :
    Expr d Γ eff₁ t → Expr d Γ eff₂ t
  | Expr.mk op ty_eq eff_le args regArgs =>
    have heff : DialectSignature.effectKind op ≤ eff₂ := by admit /- proof elided -/

content: def Expr.args {Γ ts} (e : Expr d Γ eff ts) :
    HVector (Var Γ) (DialectSignature.sig e.op) :=
  Expr.casesOn e (fun _ _ _ args _ => args)

content: def Expr.regArgs {Γ ts} (e : Expr d Γ eff ts) :
    Regions d (DialectSignature.regSig e.op) :=
  Expr.casesOn e (fun _ _ _ _ regArgs => regArgs)

content: abbrev Regions (regSig : RegionSignature d.Ty) : Type :=
  HVector (fun t => Com d t.1 .impure t.2) regSig

content: def Expr.op {Γ : Ctxt d.Ty} {eff : EffectKind} {ty} (e : Expr d Γ eff ty) : d.Op :=
  Expr.casesOn e (fun op _ _ _ _ => op)

content: def Com.denote : Com d Γ eff ty → (Γv : Valuation Γ) →
    eff.toMonad d.m (HVector toType ty)
  | .rets vs, Γv     => pure (vs.map Γv)
  | .var e body, V => e.denote V >>= body.denote

content: def Lets.denote [DialectSignature d] [DialectDenote d] {Γ₂}
    (lets : Lets d Γ₁ eff Γ₂) (V : Valuation Γ₁) : (eff.toMonad d.m <| Valuation Γ₂) :=
  match lets with
  | .nil          => return V
  | .var lets' e  => lets'.denote V >>= e.denote

content: def HVector.denote :
    {l : RegionSignature d.Ty} → (T : HVector (fun t => Com d t.1 .impure t.2) l) →
    HVector (fun t => t.1.Valuation → EffectKind.impure.toMonad d.m (HVector toType t.2)) l
  | _, .nil => HVector.nil
  | _, .cons v vs => HVector.cons (v.denote) (HVector.denote vs)

content: @[simp] abbrev FlatCom.denote [DialectDenote d]
    (flatCom : FlatCom d Γ eff Γ_out ts)
    (V : Γ.Valuation) : eff.toMonad d.m (HVector toType ts) :=
  flatCom.lets.denote V >>= (return flatCom.rets.map ·)

content: def RegionSignature.map (f : Ty → Ty') : RegionSignature Ty → RegionSignature Ty' :=
  List.map fun ⟨Γ, ty⟩ => (Γ.map f, ty.map f)

content: def Signature.map (f : Ty → Ty') : Signature Ty → Signature Ty' :=
  fun sig => {
    sig         := sig.sig.map f
    regSig      := sig.regSig.map f
    returnTypes := sig.returnTypes.map f
  }

content: def map (f : Ty₁ → Ty₂) : Ctxt Ty₁ → Ctxt Ty₂ :=
  ofList ∘ (List.map f) ∘ toList

content: def Expr.denote {ty} (e : Expr d Γ eff ty) (V : Valuation Γ) :
    eff.toMonad d.m (e.outContext.Valuation) :=
  match e with
  | ⟨op, ty_eq, heff, args, regArgs⟩ => do
      let argsDenote := args.map V
      let val ← EffectKind.liftEffect heff <| DialectDenote.denote op argsDenote regArgs.denote
      return (val ++ V).cast (by admit /- proof elided -/
      )

content: def Com.outContext {Γ} : Com d Γ eff ts → Ctxt d.Ty :=
  Com.rec' (motive := fun _ => Ctxt d.Ty)
    (@fun Γ _ => Γ) 
    (fun _ _ r => r) 

content: def Com.rec' {Γ} (com : Com d Γ eff t) : motive com :=
   
  Com.rec
    (motive_1 := fun _ _ _ _ => PUnit)
    (motive_2 := fun _ eff' t' c =>
      (h₁ : eff = eff') → (h₂ : t = t') → motive (h₁ ▸ h₂ ▸ c))
    (motive_3 := fun _ _ => PUnit)
    (fun _ _ _ _ _ _ => ⟨⟩) 
    (fun v h₁ h₂ => 
      cast (by admit /- proof elided -/
      ) <| rets (h₂ ▸ v))
    (fun e' body' _ r' h₁ h₂ => 
      let e := h₁ ▸ e'
      let body : Com _ _ eff t := cast (by admit /- proof elided -/
      ) body'
      let r : motive body := cast (by admit /- proof elided -/
      ) (r' h₁ h₂)
      cast (by admit /- proof elided -/
      ) <| var e body r)
    ⟨⟩
    (fun _ _ _ _ => ⟨⟩)
    com
    rfl
    rfl

content: def Valuation.cast {Γ Δ : Ctxt Ty} (h : Γ = Δ) (V : Valuation Γ) : Valuation Δ :=
  fun _ v => V <| v.castCtxt h.symm

content: def liftEffect [Pure m] {e1 e2 : EffectKind} {α : Type}
    (hle : e1 ≤ e2) (v1 : e1.toMonad m α) : e2.toMonad m α :=
  match e1, e2, hle with
    | .pure, .pure, _ | .impure, .impure, _ => v1
    | .pure, .impure, _ => Pure.pure v1

content: def toMonad (e : EffectKind) (m : Type → Type) : Type → Type :=
  match e with
  | pure => Id
  | impure => m

content: def Valuation.nil : Ctxt.Valuation (∅ : Ctxt Ty) := fun _ v => v.emptyElim

content: def Valuation.cons {Γ : Ctxt Ty} {t : Ty} (x : toType t) (V : Γ.Valuation) :
    (Γ.cons t).Valuation :=

content: def Com.ty : Com d Γ eff [t] → d.Ty := fun _ => t

content: infixr:50 "::ₕ" => HVector.cons

content: notation "⟦" x "⟧" => TyDenote.toType x

content: infixr:67 "::ᵥ" => Valuation.cons

-- Repository lemmas used:
content: theorem List.getElem_append_left{α : Type u_1} {l₁ l₂ : List α} {i : Nat} (hn : i < l₁.length) :\n(l₁ ++ l₂)[i] = l₁[i]

-- Local definitions used:
content: class HasTy (d : Dialect) (DenotedTy : Type) [TyDenote d.Ty] [DialectSignature d] where
    ty : d.Ty
    denote_eq : toType ty = DenotedTy := by admit /- proof elided -/

content: inductive Scf.Op (Op' Ty' : Type) (m') [TyDenote Ty'] [DialectSignature ⟨Op', Ty', m'⟩]
    [DialectDenote ⟨Op', Ty', m'⟩] : Type _
  | coe (o : Op')
  | iterate (k : ℕ) 
  | run (inputty : Ty')  
  | if (inputty retty' : Ty')  
  | for (ty : Ty')
  deriving DecidableEq, Repr

content: def Scf (d : Dialect) [TyDenote d.Ty] [DialectSignature d] [DialectDenote d] : Dialect where
  Op := Scf.Op d.Op d.Ty d.m
  Ty := d.Ty
  m  := d.m

content: abbrev LoopBody (t : Type) : Type := Int → t → t

content: def counterDecorator (δ : Int) (f : LoopBody α) : Int × α → Int × α :=
  fun (i, v) => (i + δ, f i v)

content: def IndexInvariant (f : LoopBody t) : Prop :=
  ∀ (i j : Int) (v : t), f i v = f j v

content: def atZero (f : LoopBody t) : t → t := fun v => f 0 v

content: inductive Ty
| int
| bool
| nat
 deriving DecidableEq, Repr

content: inductive Op
  | add : Op  
  | add_nat : Op  
  | axpy : Op  
  | neg : Op  
  | const : (val : ℤ) → Op
  | const_nat : (val : ℕ) → Op

content: abbrev Arith : Dialect := {Op, Ty}

content: abbrev ScfArith := Scf Arith

content: @[simp_denote] def cst (n : ℤ) : Expr ScfArith Γ .pure .int  :=
  Expr.mk
    (op := .coe <| .const n)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .nil)
    (regArgs := .nil)

content: @[simp_denote] def add {Γ : Ctxt _} (e₁ e₂ : Var Γ .int) : Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .add)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

content: @[simp_denote] def add_nat (e₁ e₂ : Var Γ .nat) : Expr ScfArith Γ .pure .nat :=
  Expr.mk
    (op := .coe <| .add_nat)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons e₁ <| .cons e₂ .nil)
    (regArgs := .nil)

content: @[simp_denote] def axpy {Γ : Ctxt _} (a : Var Γ .int) (x : Var Γ .nat) (b: Var Γ .int) :
    Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .axpy)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .cons x <| .cons b .nil)
    (regArgs := .nil)

content: @[simp_denote] def neg {Γ : Ctxt _} (a : Var Γ .int) : Expr ScfArith Γ .pure .int :=
  Expr.mk
    (op := .coe <| .neg)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons a <| .nil)
    (regArgs := .nil)

content: @[simp_denote] def iterate {Γ : Ctxt _} (k : Nat) (input : Var Γ Arith.Ty.int)
    (body : Com ScfArith ⟨[.int]⟩ .impure .int) : Expr ScfArith Γ .impure .int :=
  Expr.mk
    (op := .iterate k)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons input .nil)
    (regArgs := HVector.cons body HVector.nil)

content: @[simp_denote]
def run {Γ : Ctxt _} {t : Arith.Ty} (v : Var Γ t) (body : Com ScfArith ⟨[t]⟩ .impure t) :
    Expr ScfArith Γ .impure t :=
  Expr.mk
    (op := .run t)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons v .nil)
    (regArgs := HVector.cons body <| HVector.nil)

content: @[simp_denote] def for_ {Γ : Ctxt Arith.Ty} {t : Arith.Ty}
    (start step : Var Γ Arith.Ty.int)
    (niter : Var Γ Arith.Ty.nat) (v : Var Γ t) (body : Com ScfArith ⟨[.int, t]⟩ .impure t) :
      Expr ScfArith Γ .impure t :=
  Expr.mk
    (op := .for t)
    (ty_eq := rfl)
    (eff_le := by admit /- proof elided -/
    )
    (args := .cons start <| .cons step <| .cons niter <| .cons v .nil)
    (regArgs := HVector.cons body <| HVector.nil)

content: def lhs (vincrement : ℤ) : Com ScfArith ⟨[  .nat,   .int]⟩ .impure .int :=
    Com.letPure (cst 0) <|
    Com.letPure  (cst 1) <|
    Com.var (for_ (t := .int)
                        ⟨  1, rfl⟩
                        ⟨  0, rfl⟩
                        ⟨  2, rfl⟩
                        ⟨  3, rfl⟩ (
      Com.letPure (cst vincrement) <|
      Com.letPure (add ⟨0, rfl⟩ ⟨2, rfl⟩) 
      <| Com.ret ⟨0, rfl⟩)) <|
  Com.ret ⟨0, rfl⟩

content: def rhs (vincrement : ℤ) : Com ScfArith ⟨[  .nat,   .int]⟩ .pure .int :=
  Com.var (cst vincrement) <|
  Com.var (axpy ⟨0, rfl⟩ ⟨1, rfl⟩ ⟨2, rfl⟩) <|
  Com.ret ⟨0, rfl⟩

-- Local lemmas used:
content: theorem eq_invariant_fn
    (f : LoopBody t) (g : t → t) (hf : ∀ (i : Int) (v : t), f i v = g v) :
    LoopBody.IndexInvariant f ∧ atZero f = g

content: @[simp]
theorem eval' {f : LoopBody t} (hf : LoopBody.IndexInvariant f) (i : Int) (v : t) :
    f i v = f.atZero v

-- Target theorem and proof:
theorem correct : Com.denote (lhs v0) = Com.denote (rhs v0)
:= by
  unfold lhs rhs
  simp_peephole
  intros A B
  rw [Scf.LoopBody.counterDecorator.const_index_fn_iterate (f' := fun v => v0 + v)] <;> try rfl
  simp only [add_left_iterate, nsmul_eq_mul, Int.mul_comm]