module Data.StableTree.Types
( IsKey(..)
, Tree(..)
, Complete
, Incomplete
, Depth
, ValueCount
, nextBottom
, nextBranch
, getKey
, completeKey
, treeContents
, branchContents
, getDepth
, getValueCount
) where
import Data.StableTree.Types.Key
import qualified Data.Map as Map
import Control.Arrow ( first, second )
import Data.Map ( Map )
import Data.List ( intercalate )
data Incomplete
data Complete
type Depth = Int
type ValueCount = Int
data Tree c k v where
Bottom :: (SomeKey k, v)
-> (SomeKey k, v)
-> Map (Key Nonterminal k) v
-> (Key Terminal k, v)
-> Tree Complete k v
Branch :: Depth
-> ValueCount
-> (SomeKey k, Tree Complete k v)
-> (SomeKey k, Tree Complete k v)
-> Map (Key Nonterminal k) (Tree Complete k v)
-> (Key Terminal k, Tree Complete k v)
-> Tree Complete k v
IBottom0 :: Maybe (SomeKey k, v)
-> Tree Incomplete k v
IBottom1 :: (SomeKey k, v)
-> (SomeKey k, v)
-> Map (Key Nonterminal k) v
-> Tree Incomplete k v
IBranch0 :: Depth
-> ValueCount
-> (SomeKey k, Tree Incomplete k v)
-> Tree Incomplete k v
IBranch1 :: Depth
-> ValueCount
-> (SomeKey k, Tree Complete k v)
-> Maybe (SomeKey k, Tree Incomplete k v)
-> Tree Incomplete k v
IBranch2 :: Depth
-> ValueCount
-> (SomeKey k, Tree Complete k v)
-> (SomeKey k, Tree Complete k v)
-> Map (Key Nonterminal k) (Tree Complete k v)
-> Maybe (SomeKey k, Tree Incomplete k v)
-> Tree Incomplete k v
nextBottom :: (Ord k, IsKey k)
=> Map k v
-> Either (Tree Incomplete k v)
(Tree Complete k v, Map k v)
nextBottom values =
case Map.minViewWithKey values >>= return . second Map.minViewWithKey of
Nothing -> Left $ IBottom0 Nothing
Just ((k,v), Nothing) -> Left $ IBottom0 $ Just (wrap k, v)
Just (f1, Just (f2, remain)) ->
go (first wrap f1) (first wrap f2) Map.empty remain
where
go f1 f2 accum remain =
case Map.minViewWithKey remain of
Nothing ->
Left $ IBottom1 f1 f2 accum
Just ((k, v), remain') ->
case wrap k of
SomeKey_N nonterm ->
go f1 f2 (Map.insert nonterm v accum) remain'
SomeKey_T term ->
Right (Bottom f1 f2 accum (term, v), remain')
nextBranch :: (Ord k, IsKey k)
=> Map k (Tree Complete k v)
-> Maybe (k, Tree Incomplete k v)
-> Either (Tree Incomplete k v)
(Tree Complete k v, Map k (Tree Complete k v))
nextBranch branches mIncomplete =
let freebies = Map.minViewWithKey branches
>>= return . second Map.minViewWithKey
in case freebies of
Nothing ->
case mIncomplete of
Nothing -> Left $ IBottom0 Nothing
Just (ik, iv) -> Left $ IBranch0 depth (getValueCount iv) (wrap ik, iv)
Just ((k,v), Nothing) ->
let vcount = getValueCount v + maybe 0 (getValueCount . snd) mIncomplete
in Left $ IBranch1 depth vcount (wrap k,v) $ wrapMKey mIncomplete
Just (f1, Just (f2, remain)) ->
go (wrapKey f1) (wrapKey f2) Map.empty remain
where
go f1 f2 accum remain =
let popd = Map.minViewWithKey remain >>= return . first wrapKey
in case popd of
Nothing ->
let vcount = (getValueCount . snd) f1
+ (getValueCount . snd) f2
+ sum (map getValueCount $ Map.elems accum)
+ maybe 0 (getValueCount . snd) mIncomplete
in Left $ IBranch2 depth vcount f1 f2 accum $ wrapMKey mIncomplete
Just ((SomeKey_T term,v), remain') ->
let vcount = (getValueCount . snd) f1
+ (getValueCount . snd) f2
+ sum (map getValueCount $ Map.elems accum)
+ getValueCount v
in Right ( Branch depth vcount f1 f2 accum (term, v), remain' )
Just ((SomeKey_N nonterm,v), remain') ->
go f1 f2 (Map.insert nonterm v accum) remain'
wrapKey :: IsKey k => (k,v) -> (SomeKey k, v)
wrapKey = first wrap
wrapMKey :: IsKey k => Maybe (k,v) -> Maybe (SomeKey k, v)
wrapMKey = (>>=return . wrapKey)
depth = case Map.elems branches of
[] ->
case mIncomplete of
Nothing -> 1
Just (_, v) -> 1 + getDepth v
elems ->
let depths@(f:r) = map getDepth elems
(best, rest) = case mIncomplete of
Nothing -> (f, r)
Just (_, v) -> (getDepth v, depths)
in if all (==best) rest
then 1 + best
else error "Depth mismatch in nextBranch"
getKey :: Tree c k v -> Maybe k
getKey (Bottom (k,_) _ _ _) = Just $ unwrap k
getKey (Branch _ _ (k,_) _ _ _) = Just $ unwrap k
getKey (IBottom0 Nothing) = Nothing
getKey (IBottom0 (Just (k,_))) = Just $ unwrap k
getKey (IBottom1 (k,_) _ _) = Just $ unwrap k
getKey (IBranch0 _ _ (k,_)) = Just $ unwrap k
getKey (IBranch1 _ _ (k,_) _) = Just $ unwrap k
getKey (IBranch2 _ _ (k,_) _ _ _) = Just $ unwrap k
completeKey :: Tree Complete k v -> k
completeKey (Bottom (k,_) _ _ _) = unwrap k
completeKey (Branch _ _ (k,_) _ _ _) = unwrap k
treeContents :: Ord k => Tree c k v -> Map k v
treeContents t =
case branchContents t of
Left ( completes, Nothing) ->
Map.unions $ map treeContents $ Map.elems completes
Left ( completes, Just (_k, iv)) ->
Map.unions $ treeContents iv:map treeContents (Map.elems completes)
Right x -> x
getDepth :: Tree c k v -> Depth
getDepth (Bottom _ _ _ _) = 0
getDepth (Branch d _ _ _ _ _) = d
getDepth (IBottom0 _) = 0
getDepth (IBottom1 _ _ _) = 0
getDepth (IBranch0 d _ _) = d
getDepth (IBranch1 d _ _ _) = d
getDepth (IBranch2 d _ _ _ _ _) = d
getValueCount :: Tree c k v -> ValueCount
getValueCount (Bottom _ _ m _) = 3 + Map.size m
getValueCount (Branch _ c _ _ _ _) = c
getValueCount (IBottom0 Nothing) = 0
getValueCount (IBottom0 _) = 1
getValueCount (IBottom1 _ _ m) = 2 + Map.size m
getValueCount (IBranch0 _ c _) = c
getValueCount (IBranch1 _ c _ _) = c
getValueCount (IBranch2 _ c _ _ _ _) = c
branchContents :: Ord k
=> Tree c k v
-> Either ( Map k (Tree Complete k v)
, Maybe (k, Tree Incomplete k v))
( Map k v )
branchContents (Bottom (k1,v1) (k2,v2) terms (kt,vt)) =
let terms' = Map.mapKeys fromKey terms
conts = Map.insert (unwrap k1) v1
$ Map.insert (unwrap k2) v2
$ Map.insert (fromKey kt) vt
terms'
in Right conts
branchContents (Branch _d _c (k1,v1) (k2,v2) terms (kt,vt)) =
let terms' = Map.mapKeys fromKey terms
conts = Map.insert (unwrap k1) v1
$ Map.insert (unwrap k2) v2
$ Map.insert (fromKey kt) vt
terms'
in Left (conts, Nothing)
branchContents (IBottom0 Nothing) =
Right Map.empty
branchContents (IBottom0 (Just (k,v))) =
Right $ Map.singleton (unwrap k) v
branchContents (IBottom1 (k1,v1) (k2,v2) terms) =
let terms' = Map.mapKeys fromKey terms
conts = Map.insert (unwrap k1) v1
$ Map.insert (unwrap k2) v2
terms'
in Right conts
branchContents (IBranch0 _d _c incomplete) =
Left (Map.empty, Just $ first unwrap incomplete)
branchContents (IBranch1 _d _c (k1,v1) mIncomplete) =
Left (Map.singleton (unwrap k1) v1, mIncomplete >>= return . first unwrap)
branchContents (IBranch2 _d _c (k1,v1) (k2,v2) terms mIncomplete) =
let terms' = Map.mapKeys fromKey terms
conts = Map.insert (unwrap k1) v1
$ Map.insert (unwrap k2) v2
terms'
in Left (conts, mIncomplete >>= return . first unwrap)
instance (Ord k, Show k, Show v) => Show (Tree c k v) where
show t@(Bottom _ _ _ _) = branchShow "Bottom" t
show t@(Branch _ _ _ _ _ _) = branchShow "Branch" t
show t@(IBottom0 _) = branchShow "IBottom" t
show t@(IBottom1 _ _ _) = branchShow "IBottom" t
show t@(IBranch0 _ _ _) = branchShow "IBranch" t
show t@(IBranch1 _ _ _ _) = branchShow "IBranch" t
show t@(IBranch2 _ _ _ _ _ _) = branchShow "IBranch" t
branchShow :: (Ord k, Show k, Show v) => String -> Tree c k v -> String
branchShow header t =
case branchContents t of
Left (ts, Nothing) ->
let strs = [show k ++ " => " ++ show v | (k, v) <- Map.toAscList ts]
str = intercalate ", " strs
in header ++ "(" ++ show (getDepth t) ++ ")<" ++ str ++ ">"
Left (ts, Just (ik, iv)) ->
let strs = [ show k ++ " => " ++ show v | (k, v) <- Map.toAscList ts
] ++ [show ik ++ " => " ++ show iv]
str = intercalate ", " strs
in header ++ "(" ++ show (getDepth t) ++ ")<" ++ str ++ ">"
Right vals ->
let strs = [ show k ++ " => " ++ show v | (k, v) <- Map.toAscList vals ]
str = intercalate ", " strs
in header ++ "(" ++ show (getDepth t) ++ ")<" ++ str ++ ">"