B.hs

-- {-@ LIQUID "--reflection" @-}
{-@ LIQUID "--ple" @-}

import qualified Data.Set as Set
import           Data.Set (Set)

-- makeLeafSet = Set.fromList . map (uncurry Leaf)


data Bit = B0 | B1 deriving (Eq,Ord)

{-@ data HuffTree = 
    Leaf {symbol :: String, lw :: Nat} 
    | Tree {left :: HuffTree, right :: HuffTree, syms :: {s:Set String | s == union (symbols left) (symbols right)}, tw :: {i:Nat | i == weight left + weight right}}
@-}
data HuffTree = Leaf {symbol :: String, lw :: Int} | Tree {left :: HuffTree, right :: HuffTree, syms :: Set String, tw :: Int} deriving (Show,Eq,Ord)

{-@ measure isTree :: HuffTree -> Bool @-}
isTree Leaf{} = False
isTree Tree{} = True

{-@ measure weight :: HuffTree -> Nat @-}
weight Leaf{lw=w} = w
weight Tree{tw=w} = w

{-@ assume ww :: t:HuffTree -> {v:Nat | v == weight t} @-}
ww Leaf{lw=w} = w
ww Tree{tw=w} = w

{-@ measure symbols @-}
symbols (Leaf sym _) = Set.singleton sym
symbols (Tree _ _ syms _) = syms


instance Show Bit where
    show B0 = "0"
    show B1 = "1"

{-@ makeCodeTree :: HuffTree -> HuffTree -> HuffTree @-}
makeCodeTree :: HuffTree -> HuffTree -> HuffTree
makeCodeTree l r = Tree l r (Set.union (symbols l) (symbols r)) (ww l + ww r)

{-@ chooseBranch :: Bit -> {t:HuffTree | isTree t} -> HuffTree @-}
chooseBranch B0 Tree{left=l,right=r} = l
chooseBranch B1 Tree{left=l,right=r} = r
chooseBranch b t = error "no"

{-@ chooseBranchSymbol :: sym:String -> {t:HuffTree | isTree t && member sym (symbols t)} -> (HuffTree, Bit) @-}
chooseBranchSymbol sym t@Tree{left=l,right=r} 
    | Set.member sym (symbols l) = (l,B0)
    | Set.member sym (symbols r) = (r,B1)
    -- | otherwise = error $ show (sym,t)
chooseBranchSymbol sym t = error "no" 

{-@ measure notNull @-}
notNull [] = False
notNull (_:_) = True

-- {-@ encode :: syms:[String] -> {t:HuffTree | notNull syms => isTree t} -> [Bit] @-}
{-@ encode :: syms:[String] -> {t:HuffTree | True} -> [Bit] @-}
encode syms tree = encode' syms tree
    where
        encode' [] curBranch = []
        encode' (s:ss) curBranch =
            case chooseBranchSymbol s curBranch of
                (Leaf{},b)      -> b : encode' ss tree
                (t@Tree{},b)    -> b : encode' (s:ss) t
                



{-@ decode :: [Bit] -> HuffTree -> [String] @-}
decode bits tree = decode' bits tree
    where
        decode' [] cur = []
        decode' (s:ss) cur = 
            case chooseBranch s cur of
                Leaf{symbol=s} -> s : decode' ss tree
                t@Tree{} -> decode' ss t

sampleCode = [B0,B1,B0,B0,B1,B1,B1,B1,B1,B0]
sampleTree = 
  makeCodeTree 
   (Leaf "ninjas" 8) 
   (makeCodeTree 
    (Leaf "fight" 5) 
    (makeCodeTree 
     (Leaf "night" 1) 
     (Leaf "by" 1)))

main = do
    print $ decode (encode ["ninjas","fight","ninjas","by","night"] sampleTree) sampleTree