// 応用編のコードを全部くっつけたコード

// コード5: マークルツリーの定義
module MerkleTree
    // 長さnの文字列の型lstring nを定義
    let lstring (n:nat) = 
        s:string{String.length s == n}

    // 文字列を連結させる関数concatを定義
    let concat #n #m (s0:lstring n) (s1:lstring m) : lstring (m + n)
        = String.concat_length s0 s1; 
            s0 ^ s1

    // -ハッシュ関数の出力サイズ hash_size
    // -ハッシュ関数の出力の型 hash_t
    // -ハッシュ関数 hash
    // -ツリーに格納するデータの型 mtdata
    // を定義
    assume val hash_size:nat
    let hash_t = lstring hash_size
    assume val hash (m:string) : hash_t
    let mtdata = string

    type mtree: nat -> hash_t -> Type =
    // データのハッシュ値が格納されている高さ0のマークルツリーとして葉ノードLを定義
    | L: 
        data:mtdata ->
        mtree 0 (hash data)
    // 左右のツリーから内部ノードNを定義．格納されている値は左右のツリーのルートノードの連結のハッシュ値
    | N:
        #n:nat ->
        #hl:hash_t ->
        #hr:hash_t ->
        left:mtree n hl ->
        right:mtree n hr ->
        mtree (n + 1) (hash (concat hl hr))


// コード6: データアクセス関数の定義
module MTAccess
    open MerkleTree
    module Lib = FStar.List.Tot
  
    // ブール値のリストでデータIDの型data_idを定義
    let data_id = list bool

    // データIDとマークルツリーを受け取りIDに対応するデータを返す関数getを定義
    let rec get #h
            (did:data_id)
            (tree:mtree (Lib.length did) h)
    : Tot mtdata (decreases did)
    // データIDを1ビットずつ読み込み，1ならば左側のサブツリーに0ならば右側のサブツリーに移動し，葉ノードに到達したらそこに格納されているデータを返す
    = match did with
    | [] -> L?.data tree
    | b::did' ->
        if b then get did' (N?.left tree)
        else      get did' (N?.right tree)
        
        
// コード7: 証拠付きデータアクセス関数の定義
module MTAccessEvidence
    open MerkleTree
    open MTAccess
    module Lib = FStar.List.Tot

    type data_with_evidence : nat -> Type =
    | DATA:
        data:mtdata ->
        did:data_id ->
        hashes:list hash_t {Lib.length did == Lib.length hashes} ->
        data_with_evidence (Lib.length did)

    let rec get_with_evidence (#h:_) (did:data_id) 
                                (tree:mtree (Lib.length did) h)
    :Tot (data_with_evidence (Lib.length did)) (decreases did)
    = match did with
    | [] ->
        DATA (L?.data tree) [] []
    | b::did' ->
        let N #_ #hl #hr left right = tree in 
        if b then
            let p = get_with_evidence did' left in
            DATA p.data did (hr :: p.hashes)
        else
            let p = get_with_evidence did' right in
            DATA p.data did (hl :: p.hashes)


// コード8: 証拠を検証する関数の定義
module VerifyEvidence
    open MerkleTree
    open MTAccess
    open MTAccessEvidence
    module Lib = FStar.List.Tot

    // 長さnの証拠付きデータからそれぞれの要素の後ろ部分を抽出した長さn-1の証拠付きデータを返す関数
    let tail #n (p:data_with_evidence n{n>0})
        : data_with_evidence (n-1)
        = DATA p.data (Lib.tail p.did) (Lib.tail p.hashes)

    let rec compute_root_hash (#n:nat) (p:data_with_evidence n)
        : hash_t
        // はじめにデータのハッシュ値を計算し，それ以降はデータIDの後ろ方向からハッシュを再計算
        // データIDに基づき，左側または右側の兄弟ハッシュを連結してそのハッシュ値を計算
        = let DATA d did hashes = p in
        match did with
        | [] -> hash p.data
        | bit::did' ->
            let h' = compute_root_hash (tail p) in
            if bit then hash (concat h' (Lib.hd hashes))
            else hash (concat (Lib.hd hashes) h')
    
    // 証拠付きデータとマークルツリーを受け取り，証拠から計算されたハッシュ値とマークルツリーのルートハッシュが一致すればtrueを，そうでなければfalseを返す関数
    let verify #h #n (p:data_with_evidence n) (tree:mtree n h)
        : bool
        = compute_root_hash p = h

    // get_with_evidence関数が返す証拠はいつでも正しい（verify関数でtrueと評価される）ことを証明
    let rec correctness (#h:hash_t) (did:data_id)
                      (tree:mtree (Lib.length did) h)
        : Lemma (ensures (verify (get_with_evidence did tree) tree))
                (decreases did)
        = match did with
        | [] -> ()
        | bit::did' ->
            let N left right = tree in
            if bit then
                correctness did' left
            else
                correctness did' right
                

// コード9: 安全性の証明
module ProofSecurity
    open MerkleTree
    open MTAccess
    open MTAccessEvidence
    open VerifyEvidence
    module Lib = FStar.List.Tot

    // ハッシュ値が同じ値になる2つの異なる文字列のペアの型を定義
    type hash_collision =
    | Collision :
        s1:string ->
        s2:string {hash s1 = hash s2 /\ not (s1 = s2)} ->
        hash_collision
    
    // verify関数でtrueと評価される不正な証拠からhash_collision型の値を返す関数を作れることを示すことで安全性を証明
    let rec security (#n:nat) (#h:hash_t) (tree:mtree n h)
                   (p:data_with_evidence n {verify p tree/\not (get p.did tree = p.data)})
    : hash_collision
    = match p.did with
    | [] -> Collision p.data (L?.data tree)
    | b::did' ->
        let N #_ #h1 #h2 left right = tree in
        let h' = compute_root_hash (tail p) in
        let hd :: _ = p.hashes in
            if b then
                if h' = h1 then security left (tail p)
                else (String.concat_injective h1 h' h2 hd;
                      Collision (concat h1 h2) (concat h' hd))
            else
                if h' = h2 then security right (tail p)
                else (String.concat_injective h1 hd h2 h';
                      Collision (concat h1 h2) (concat hd h'))