------------------------------------------------------------------------
-- Programs used to define and specify breadth-first manipulations of
-- trees
------------------------------------------------------------------------

module BreadthFirst.Programs where

open import Codata.Musical.Notation
open import Codata.Musical.Colist hiding ([_])
open import Data.List.NonEmpty using (List⁺; [_]; _⁺++⁺_)
open import Data.Product

open import BreadthFirst.Universe
open import Stream hiding (zipWith)
open import Tree

infixr 5 _≺_ _∷_
infixr 4 _,_
infix  3 ↓_

mutual

  -- The term WHNF is a bit of a misnomer here; only recursive
  -- /coinductive/ arguments are suspended (in the form of programs).

  data ElW : U  Set₁ where
    leaf :  {a}  ElW (tree a)
    node :  {a}
           (l :  (ElP (tree a))) (x : ElW a) (r :  (ElP (tree a))) 
           ElW (tree a)
    _≺_  :  {a} (x : ElW a) (xs :  (ElP (stream a)))  ElW (stream a)
    []   :  {a}  ElW (colist a)
    _∷_  :  {a} (x : ElW a) (xs :  (ElP (colist a)))  ElW (colist a)
    _,_  :  {a b} (x : ElW a) (y : ElW b)  ElW (a  b)
    ⌈_⌉  :  {A} (x : A)  ElW  A 

  -- Note that higher-order functions are not handled well (see
  -- longZipWith): the arguments cannot be programs, since this would
  -- make the type negative.

  -- The tree arguments to lab and flatten used to be tree programs,
  -- but a number of lemmas could be removed when they were turned
  -- into "actual" trees.

  data ElP : U  Set₁ where
    ↓_          :  {a} (w : ElW a)  ElP a
    fst         :  {a b} (p : ElP (a  b))  ElP a
    snd         :  {a b} (p : ElP (a  b))  ElP b
    lab         :  {A B}
                  (t : Tree A) (lss : ElP (stream  Stream B )) 
                  ElP (tree  B   stream  Stream B )
    longZipWith :  {A} (f : A  A  A) (xs ys : ElP (colist  A )) 
                  ElP (colist  A )
    flatten     :  {A} (t : Tree A)  ElP (colist  List⁺ A )

infixl 9 _·_ _⊛_

_·_ :  {A B}  (A  B)  ElW  A   ElW  B 
f ·  x  =  f x 

_⊛_ :  {A B}  ElW  (A  B)   ElW  A   ElW  B 
 f   x = f · x

fstW :  {a b}  ElW (a  b)  ElW a
fstW (x , y) = x

sndW :  {a b}  ElW (a  b)  ElW b
sndW (x , y) = y

-- Uses the n-th stream to label the n-th level in the tree. Returns
-- the remaining stream elements (for every level).

labW :  {A B}  Tree A  ElW (stream  Stream B ) 
       ElW (tree  B   stream  Stream B )
labW leaf         bss                = (leaf , bss)
labW (node l _ r) ( b  bs   bss) =
  (node ( fst x)  b  ( fst y) ,   bs    snd y)
  where
  x = lab ( l) ( bss)
  y = lab ( r) (snd x)

longZipWithW :  {A}  (A  A  A) 
               (xs ys : ElW (colist  A ))  ElW (colist  A )
longZipWithW f (x  xs) (y  ys) = f · x  y   longZipWith f ( xs) ( ys)
longZipWithW f xs       []       = xs
longZipWithW f []       ys       = ys

flattenW :  {A}  Tree A  ElW (colist  List⁺ A )
flattenW leaf         = []
flattenW (node l x r) =
   [ x ]    longZipWith _⁺++⁺_ (flatten ( l)) (flatten ( r))

whnf :  {a}  ElP a  ElW a
whnf ( w)                 = w
whnf (fst p)               = fstW (whnf p)
whnf (snd p)               = sndW (whnf p)
whnf (lab t bss)           = labW t (whnf bss)
whnf (longZipWith f xs ys) = longZipWithW f (whnf xs) (whnf ys)
whnf (flatten t)           = flattenW t

mutual

  ⟦_⟧W :  {a}  ElW a  El a
   leaf         ⟧W = leaf
   (node l x r) ⟧W = node (   l )  x ⟧W (   r )
   (x  xs)     ⟧W =  x ⟧W     xs 
   []           ⟧W = []
   (x  xs)     ⟧W =  x ⟧W     xs 
   (x , y)      ⟧W = ( x ⟧W ,  y ⟧W)
    x         ⟧W = x

  ⟦_⟧ :  {a}  ElP a  El a
   p  =  whnf p ⟧W