NestedAndFlattenedTrees.agda

open import Data.List
open import Data.Nat
open import Data.Product
open import Data.Vec renaming (_++_ to _++V_ ; replicate to replicateV)
open import Relation.Binary.PropositionalEquality



data Tree (A : Set) : Set where
  leaf : A → Tree A
  branch : (n : ℕ) → Vec (Tree A) n → Tree A


data LinTreePart (A : Set) : List Set → List Set → Set₁ where
  leaf : ∀ {S} → A → LinTreePart A S (Tree A ∷ S)
  branch : ∀ {S} → (n : ℕ) → LinTreePart A (replicate n (Tree A) ++ S) (Tree A ∷ S)

data LinTree (A : Set) : List Set → Set₁ where
  [] : LinTree A []
  _∷_ : ∀ {S S'} → LinTreePart A S S' → LinTree A S → LinTree A S'

mutual
  fwdOnto : ∀ {A S} → Tree A → LinTree A S → LinTree A (Tree A ∷ S)
  fwdOnto (leaf x) lt = leaf x ∷ lt
  fwdOnto (branch n ts) lt = branch n ∷ unrollOnto ts lt
  
  unrollOnto : ∀ {A S n} → (ts : Vec (Tree A) n) → LinTree A S → LinTree A (replicate n (Tree A) ++ S)
  unrollOnto [] lt = lt
  unrollOnto (t ∷ ts) lt = fwdOnto t (unrollOnto ts lt)


fwd : ∀ {A} → Tree A → LinTree A (Tree A ∷ [])
fwd t = fwdOnto t []

{-# TERMINATING #-}
mutual
  bwdOff : ∀ {A S} → LinTree A (Tree A ∷ S) → Tree A × LinTree A S
  bwdOff (leaf x ∷ lt) = leaf x , lt
  bwdOff (branch n ∷ lt) with rollOff n lt
  ... | ts , lt' = branch n ts , lt'
  
  rollOff : ∀ {A S} → (n : ℕ) → LinTree A (replicate n (Tree A) ++ S) → Vec (Tree A) n × LinTree A S
  rollOff zero lt = [] , lt
  rollOff (suc n) lt with bwdOff lt
  ... | (t , lt') with rollOff n lt'
  ... | (ts , lt'') = (t ∷ ts) , lt''

bwd : ∀ {A} → LinTree A (Tree A ∷ []) → Tree A
bwd lt = proj₁ (bwdOff lt)

mutual
  bwdOff-fwdOnto-id : ∀ {A S} → (t : Tree A) (lt : LinTree A S) → bwdOff (fwdOnto t lt) ≡ (t , lt)
  bwdOff-fwdOnto-id (leaf x) lt = refl
  bwdOff-fwdOnto-id (branch n ts) lt rewrite rollOff-unrollOnto-id ts lt = refl

  rollOff-unrollOnto-id : ∀ {A S n} → (ts : Vec (Tree A) n) (lt : LinTree A S) → rollOff n (unrollOnto ts lt) ≡ (ts , lt)
  rollOff-unrollOnto-id [] lt = refl
  rollOff-unrollOnto-id (t ∷ ts) lt
    rewrite bwdOff-fwdOnto-id t (unrollOnto ts lt)
          | rollOff-unrollOnto-id ts lt
          = refl

bwd-fwd : ∀ {A} → (t : Tree A) → bwd (fwd t) ≡ t
bwd-fwd t = cong proj₁ (bwdOff-fwdOnto-id t [])

{-# TERMINATING #-}
mutual
  fwdOnto-bwdOff-id : ∀ {A S} → (lt : LinTree A (Tree A ∷ S)) → uncurry fwdOnto (bwdOff lt) ≡ lt
  fwdOnto-bwdOff-id (leaf x ∷ lt) = refl
  fwdOnto-bwdOff-id (branch n ∷ lt) rewrite unrollOnto-rollOff-id n lt = refl

  unrollOnto-rollOff-id : ∀ {A S} → (n : ℕ) → (lt : LinTree A (replicate n (Tree A) ++ S)) → uncurry unrollOnto (rollOff n lt) ≡ lt
  unrollOnto-rollOff-id zero lt = refl
  unrollOnto-rollOff-id (suc n) lt with bwdOff lt | inspect bwdOff lt | fwdOnto-bwdOff-id lt
  ... | (t , lt') | [ eq ] | eq' rewrite unrollOnto-rollOff-id n lt' = eq'

bwdOff-empty : ∀ {A} → (lt : LinTree A (Tree A ∷ [])) → proj₂ (bwdOff lt) ≡ []
bwdOff-empty (leaf x ∷ []) = refl
bwdOff-empty (leaf x ∷ (() ∷ lt))
bwdOff-empty (branch n ∷ lt) with rollOff n lt
bwdOff-empty (branch n ∷ lt) | t , [] = refl
bwdOff-empty (branch n ∷ lt) | t , (() ∷ lt')

fwd-bwd : ∀ {A} → (lt : LinTree A (Tree A ∷ [])) → fwd (bwd lt) ≡ lt
fwd-bwd lt rewrite sym (bwdOff-empty lt) = fwdOnto-bwdOff-id lt