-- ArkLib/CommitmentScheme/InductiveMerkleTree.lean

import VCVio

import ArkLib.ToMathlib.Data.IndexedBinaryTree.Basic

import ArkLib.CommitmentScheme.Basic

import Mathlib.Data.Vector.Snoc

import ArkLib.ToVCVio.Oracle

namespace InductiveMerkleTree

open List OracleSpec OracleComp BinaryTree

section spec

variable (α : Type)

end spec

variable {α : Type}

def buildMerkleTree_with_hash {s} (leaf_tree : LeafData α s) (hashFn : α → α → α) :
    (FullData α s) :=
  match leaf_tree with
  | LeafData.leaf a => FullData.leaf a
  | LeafData.internal left right =>
    let leftTree := buildMerkleTree_with_hash left hashFn
    let rightTree := buildMerkleTree_with_hash right hashFn
    let rootHash := hashFn (leftTree.getRootValue) (rightTree.getRootValue)
    FullData.internal rootHash leftTree rightTree

def generateProof {s} (cache_tree : FullData α s) :
    BinaryTree.SkeletonLeafIndex s → List α
  | .ofLeaf => []
  | .ofLeft idxLeft =>
    (cache_tree.rightSubtree).getRootValue ::
      (generateProof cache_tree.leftSubtree idxLeft)
  | .ofRight idxRight =>
    (cache_tree.leftSubtree).getRootValue ::
      (generateProof cache_tree.rightSubtree idxRight)

def getPutativeRoot_with_hash {s} (idx : BinaryTree.SkeletonLeafIndex s)
    (leafValue : α) (proof : List α) (hashFn : α → α → α) : α :=
  match proof with
  | [] => leafValue 
  | siblingBelowRootHash :: restProof =>
    match idx with
    | BinaryTree.SkeletonLeafIndex.ofLeaf =>
      
      
      
      leafValue
    | BinaryTree.SkeletonLeafIndex.ofLeft idxLeft =>
      
      hashFn (getPutativeRoot_with_hash idxLeft leafValue restProof hashFn) siblingBelowRootHash
    | BinaryTree.SkeletonLeafIndex.ofRight idxRight =>
      
      hashFn siblingBelowRootHash (getPutativeRoot_with_hash idxRight leafValue restProof hashFn)

-- Repository definitions used:
content: def FullData.leftSubtree {α : Type} {s_left s_right : Skeleton}
    (tree : FullData α (Skeleton.internal s_left s_right)) :
    FullData α s_left :=
  match tree with
  | FullData.internal _ left _right =>
    left

content: inductive Skeleton where
  | leaf : Skeleton
  | internal : Skeleton → Skeleton → Skeleton

content: inductive FullData (α : Type) : Skeleton → Type
  | leaf (value : α) : FullData α Skeleton.leaf
  | internal {left right} (value : α)
      (leftData : FullData α left)
      (rightData : FullData α right) :
      FullData α (Skeleton.internal left right)
  deriving Repr

content: inductive Skeleton :
    (n : ℕ) → (ar : Fin n → ℕ) → Type where
  | leaf {ar0 : Fin 0 → ℕ} : Skeleton 0 ar0
  | node {n : ℕ} {ar : Fin (n+1) → ℕ}
      (children : Fin (ar 0) →
        Skeleton n (fun i => ar i.succ)) :
      Skeleton (n+1) ar

content: def FullData.rightSubtree {α : Type} {s_left s_right : Skeleton}
    (tree : FullData α (Skeleton.internal s_left s_right)) :
    FullData α s_right :=
  match tree with
  | FullData.internal _ _left right =>
    right

content: def LeafData.rightSubtree {α : Type} {s_left s_right : Skeleton}
    (tree : LeafData α (Skeleton.internal s_left s_right)) :
    LeafData α s_right :=
  match tree with
  | LeafData.internal _left right =>
    right

content: inductive LeafData (α : Type) : Skeleton → Type
  | leaf (value : α) : LeafData α Skeleton.leaf
  | internal {left right} (leftData : LeafData α left) (rightData : LeafData α right) :
      LeafData α (Skeleton.internal left right)
  deriving Repr

content: inductive SkeletonNodeIndex : Skeleton → Type
  | ofLeaf : SkeletonNodeIndex Skeleton.leaf
  | ofInternal {left right} :
      SkeletonNodeIndex (Skeleton.internal left right)
  | ofLeft {left right : Skeleton} (idxLeft : SkeletonNodeIndex left) :
      SkeletonNodeIndex (Skeleton.internal left right)
  | ofRight {left right : Skeleton} (idxRight : SkeletonNodeIndex right) :
      SkeletonNodeIndex (Skeleton.internal left right)

content: def InternalData.rightSubtree {α : Type} {s_left s_right : Skeleton}
    (tree : InternalData α (Skeleton.internal s_left s_right)) :
    InternalData α s_right :=
  match tree with
  | InternalData.internal _ _left right =>
    right

content: inductive InternalData (α : Type) : Skeleton → Type
  | leaf : InternalData α Skeleton.leaf
  | internal {left right} (value : α)
      (leftData : InternalData α left)
      (rightData : InternalData α right) : InternalData α (Skeleton.internal left right)
  deriving Repr

content: def InternalData.leftSubtree {α : Type} {s_left s_right : Skeleton}
    (tree : InternalData α (Skeleton.internal s_left s_right)) :
    InternalData α s_left :=
  match tree with
  | InternalData.internal _ left _right =>
    left

content: inductive SkeletonLeafIndex : Skeleton → Type
  | ofLeaf : SkeletonLeafIndex Skeleton.leaf
  | ofLeft {left right : Skeleton} (idxLeft : SkeletonLeafIndex left) :
      SkeletonLeafIndex (Skeleton.internal left right)
  | ofRight {left right : Skeleton} (idxRight : SkeletonLeafIndex right) :
      SkeletonLeafIndex (Skeleton.internal left right)

content: def LeafData.leftSubtree {α : Type} {s_left s_right : Skeleton}
    (tree : LeafData α (Skeleton.internal s_left s_right)) :
    LeafData α s_left :=
  match tree with
  | LeafData.internal left _right =>
    left

content: def FullData.getRootValue {s} {α : Type} (tree : FullData α s) :=
  tree.get (getRootIndex s)

content: def FullData.get {s} {α : Type}
    (tree : FullData α s) (idx : SkeletonNodeIndex s) : α :=
  match tree, idx with
  | FullData.leaf value, SkeletonNodeIndex.ofLeaf => value
  | FullData.internal value _ _, SkeletonNodeIndex.ofInternal => value
  | FullData.internal _ left _, SkeletonNodeIndex.ofLeft idxLeft =>
    FullData.get left idxLeft
  | FullData.internal _ _ right, SkeletonNodeIndex.ofRight idxRight =>
    FullData.get right idxRight

content: def InternalData.get {s} {α : Type}
    (tree : InternalData α s) (idx : SkeletonInternalIndex s) : α :=
  match tree, idx with
  | InternalData.internal value _ _, SkeletonInternalIndex.ofInternal => value
  | InternalData.internal _ left _, SkeletonInternalIndex.ofLeft idxLeft =>
    InternalData.get left idxLeft
  | InternalData.internal _ _ right, SkeletonInternalIndex.ofRight idxRight =>
    InternalData.get right idxRight

content: inductive SkeletonInternalIndex : Skeleton → Type
  | ofInternal {left right} : SkeletonInternalIndex (Skeleton.internal left right)
  | ofLeft {left right : Skeleton} (idxLeft : SkeletonInternalIndex left) :
      SkeletonInternalIndex (Skeleton.internal left right)
  | ofRight {left right : Skeleton} (idxRight : SkeletonInternalIndex right) :
      SkeletonInternalIndex (Skeleton.internal left right)

content: def LeafData.get {s} {α : Type}
    (tree : LeafData α s) (idx : SkeletonLeafIndex s) : α :=
  match tree, idx with
  | LeafData.leaf value, SkeletonLeafIndex.ofLeaf => value
  | LeafData.internal left _, SkeletonLeafIndex.ofLeft idxLeft =>
    LeafData.get left idxLeft
  | LeafData.internal _ right, SkeletonLeafIndex.ofRight idxRight =>
    LeafData.get right idxRight

content: def getRootIndex (s : Skeleton) : SkeletonNodeIndex s := match s with
  | Skeleton.leaf => SkeletonNodeIndex.ofLeaf
  | Skeleton.internal _ _ =>
    SkeletonNodeIndex.ofInternal

-- Repository lemmas used:
content: @[simp]
theorem LeafData.rightSubtree_internal {α} {s_left s_right : Skeleton}
    (left : LeafData α s_left) (right : LeafData α s_right) :
    (LeafData.internal left right).rightSubtree = right

content: @[simp]
theorem LeafData.leftSubtree_internal {α} {s_left s_right : Skeleton}
    (left : LeafData α s_left) (right : LeafData α s_right) :
    (LeafData.internal left right).leftSubtree = left

content: @[simp]
theorem LeafData.get_ofLeft {s_left s_right : Skeleton} {α : Type}
    (tree : LeafData α (Skeleton.internal s_left s_right))
    (idxLeft : SkeletonLeafIndex s_left) :
    tree.get (SkeletonLeafIndex.ofLeft idxLeft) =
      tree.leftSubtree.get idxLeft

content: @[simp]
theorem LeafData.get_ofRight {s_left s_right : Skeleton} {α : Type}
    (tree : LeafData α (Skeleton.internal s_left s_right))
    (idxRight : SkeletonLeafIndex s_right) :
    tree.get (SkeletonLeafIndex.ofRight idxRight) =
      tree.rightSubtree.get idxRight

content: @[simp]
theorem FullData.internal_getRootValue {s_left s_right : Skeleton} {α : Type}
    (value : α) (left : FullData α s_left) (right : FullData α s_right) :
    (FullData.internal value left right).getRootValue =
      value

-- Local definitions used:
content: def buildMerkleTree_with_hash {s} (leaf_tree : LeafData α s) (hashFn : α → α → α) :
    (FullData α s) :=
  match leaf_tree with
  | LeafData.leaf a => FullData.leaf a
  | LeafData.internal left right =>
    let leftTree := buildMerkleTree_with_hash left hashFn
    let rightTree := buildMerkleTree_with_hash right hashFn
    let rootHash := hashFn (leftTree.getRootValue) (rightTree.getRootValue)
    FullData.internal rootHash leftTree rightTree

content: def generateProof {s} (cache_tree : FullData α s) :
    BinaryTree.SkeletonLeafIndex s → List α
  | .ofLeaf => []
  | .ofLeft idxLeft =>
    (cache_tree.rightSubtree).getRootValue ::
      (generateProof cache_tree.leftSubtree idxLeft)
  | .ofRight idxRight =>
    (cache_tree.leftSubtree).getRootValue ::
      (generateProof cache_tree.rightSubtree idxRight)

content: def getPutativeRoot_with_hash {s} (idx : BinaryTree.SkeletonLeafIndex s)
    (leafValue : α) (proof : List α) (hashFn : α → α → α) : α :=
  match proof with
  | [] => leafValue 
  | siblingBelowRootHash :: restProof =>
    match idx with
    | BinaryTree.SkeletonLeafIndex.ofLeaf =>
      
      
      
      leafValue
    | BinaryTree.SkeletonLeafIndex.ofLeft idxLeft =>
      
      hashFn (getPutativeRoot_with_hash idxLeft leafValue restProof hashFn) siblingBelowRootHash
    | BinaryTree.SkeletonLeafIndex.ofRight idxRight =>
      
      hashFn siblingBelowRootHash (getPutativeRoot_with_hash idxRight leafValue restProof hashFn)

-- Local lemmas used:
content: @[simp]
theorem generateProof_ofLeft {sleft sright : Skeleton}
    (cache_tree : FullData α (Skeleton.internal sleft sright))
    (idxLeft : SkeletonLeafIndex sleft) :
    generateProof cache_tree (BinaryTree.SkeletonLeafIndex.ofLeft idxLeft) =
      (cache_tree.rightSubtree).getRootValue ::
        (generateProof cache_tree.leftSubtree idxLeft)

content: @[simp]
theorem generateProof_ofRight {sleft sright : Skeleton}
    (cache_tree : FullData α (Skeleton.internal sleft sright))
    (idxRight : SkeletonLeafIndex sright) :
    generateProof cache_tree (BinaryTree.SkeletonLeafIndex.ofRight idxRight) =
      (cache_tree.leftSubtree).getRootValue ::
        (generateProof cache_tree.rightSubtree idxRight)

-- Target theorem and proof:
theorem functional_completeness (α : Type) {s : Skeleton}
  (idx : SkeletonLeafIndex s)
  (leaf_data_tree : LeafData α s)
  (hash : α → α → α) :
  (getPutativeRoot_with_hash
    idx
    (leaf_data_tree.get idx)
    (generateProof
      (buildMerkleTree_with_hash leaf_data_tree hash) idx)
    (hash)) =
  (buildMerkleTree_with_hash leaf_data_tree hash).getRootValue
:= by
  induction s with
  | leaf =>
    match leaf_data_tree with
    | LeafData.leaf a =>
      cases idx with
      | ofLeaf =>
        simp [buildMerkleTree_with_hash, getPutativeRoot_with_hash]
  | internal s_left s_right left_ih right_ih =>
    match leaf_data_tree with
    | LeafData.internal left right =>
      cases idx with
      | ofLeft idxLeft =>
        simp_rw [LeafData.get_ofLeft, LeafData.leftSubtree_internal, buildMerkleTree_with_hash,
          generateProof_ofLeft, FullData.rightSubtree, FullData.leftSubtree,
          getPutativeRoot_with_hash, left_ih, FullData.internal_getRootValue]
      | ofRight idxRight =>
        simp_rw [LeafData.get_ofRight, LeafData.rightSubtree_internal, buildMerkleTree_with_hash,
          generateProof_ofRight, FullData.leftSubtree, FullData.rightSubtree,
          getPutativeRoot_with_hash, right_ih, FullData.internal_getRootValue]