content: {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE FlexibleContexts #-}
  module KMC.Determinization where
  
  import           Control.Applicative
  import           Control.Monad (guard, when)
  import           Control.Monad.State
  import           Data.Maybe (maybeToList, isJust)
  import           Data.Monoid
  import qualified Data.Set as S
  import           KMC.SymbolicFST
  import           KMC.SymbolicSST
  import           KMC.Theories
  import           KMC.TreeWriter
  import           KMC.Util.List (lcpMany1, linv')
  import           Prelude
  
  -- | The functions in the symbolic transitions generated by the path tree
  -- determinization procedure are obtained by extending the functions in the
  -- determinized FST. Namely, by adjoining an epsilon element to the set of
  -- functions and extending functions to sequences.
  type SSTFunc f = [EpsFunc f]
  
  -- | A variable is identified by its position in the path tree. The position is
  -- represented as a sequence of integers representing the path from the root of
  -- the path tree to a given node.
  data Var = Var [Int]
    deriving (Eq, Ord, Show)
  
  instance PartialOrder Var where
    lte (Var l) (Var r) = reverse l `prefixOf` reverse r
        where
          prefixOf [] _ = True
          prefixOf (x:xs) (y:ys) = (x==y) && prefixOf xs ys
          prefixOf _ _ = False
  
  
  ----------------
  -- Closure monad
  ----------------
  
  -- | The closure monad models a computation consisting of concurrent threads
  -- with shared state executed in round-robin
  type Closure st w = TreeWriterT w (State (S.Set st))
  type PathTree w a = Maybe (Tree w a)
  
  visit :: (Monoid w, Ord st) => st -> Closure st w st
  visit st = do { vis <- get;
                  when (S.member st vis) zero;
                  modify (S.insert st);
                  return st
                }
  
  -- | Resume a closure computation
  resume :: (Monoid w) => Tree w a -> Closure st w a
  resume (Tip w x) = tell w >> return x
  resume (Fork w ts) = tell w >> branch (map resume ts)
  
  -- | Resume a possibly failing closure computation
  resume' :: Monoid w => PathTree w a -> Closure st w a
  resume' Nothing = zero
  resume' (Just t) = resume t
  
  -- | Reduce a tree by removing longest common prefixes
  reduce :: (Eq a) => Tree [a] b -> Tree [a] b
  reduce (Tip w b)   = Tip w b
  reduce (Fork w ts) =
    Fork (w ++ p) $ [ t { tOutput = linv' p (tOutput t) } | t <- ts' ]
    where
      ts' = map reduce ts
      p   = lcpMany1 (map tOutput ts')
  
  -- | Interpret a closure computation as a path tree and reduce the result.
  runClosure :: (Eq gamma) => Closure st [gamma] a -> PathTree [gamma] a
  runClosure x = reduce <$> evalState (runTreeWriterT x) S.empty
  
  
  -----------------------
  -- Computing path trees
  -----------------------
  
  genclosure :: (Ord st, Monoid w)
             => FST st pred func
             -> (Rng func -> w)
             -> st
             -> Closure st w st
  genclosure fst' inj q =
    case fstEvalEpsilonEdges fst' q of
      [] -> return q
      qs -> branch $
              map (\(out, q') -> tell (inj out)
                                 >> visit q'
                                 >> genclosure fst' inj q') qs
  
  -- | Non-deterministically follow all non-input transitions.
  closure :: (Ord st)
             => FST st pred func
             -> st
             -> Closure st (UpdateString var (Rng func)) st
  closure fst' q = genclosure fst' (\c -> [Right c]) q
  
  closureFunc :: (Ord st) =>
                 FST st pred func
              -> st
              -> Closure st (UpdateStringFunc var (SSTFunc func)) st
  closureFunc fst' q = genclosure fst' (\c -> [ConstA c]) q
  
  consume :: (PartialOrder pred, Ord st)
          => (func -> func') -> FST st pred func -> pred -> st -> Closure st (UpdateStringFunc var func') st
  consume inj fst' p q =
    case fstAbstractEvalEdgesAll fst' q p of
      []        -> zero
      [(f, q')] -> tell [FuncA (inj f)] >> visit q'
      _ -> error "Stepping for FSTs with read-fanout greater than one is not supported yet"
  
  eof :: (Monoid w, Ord st) => FST st pred func -> st -> Closure st w st
  eof fst' q | S.member q (fstF fst') = visit q
             | otherwise = zero
  
  kill :: (Ord st) => S.Set st -> st -> Closure st (UpdateStringFunc var func) st
  kill kills q | S.member q kills = zero
               | otherwise = return q
  
  closureAbstractTree :: (Eq var, Eq (Rng func), Ord a) =>
                         FST a pred func
                      -> PathTree var a
                      -> PathTree (UpdateString var (Rng func)) a
  closureAbstractTree fst' tr =
    let tr' = fmap (mapOutput (return . Left)) tr
    in runClosure (resume' tr' >>= closure fst')
  
  closureTree :: (Eq var, Eq (Rng func), Ord st) =>
                 FST st pred func
              -> PathTree (UpdateString var (Rng func)) st
              -> PathTree (UpdateString var (Rng func)) st
  closureTree fst' tr = runClosure (resume' tr >>= closure fst')
  
  closureTreeFunc :: (Eq var, Eq func, Eq (Rng func), Ord st) =>
                 FST st pred func
              -> PathTree (UpdateStringFunc var (SSTFunc func)) st
              -> PathTree (UpdateStringFunc var (SSTFunc func)) st
  closureTreeFunc fst' tr = runClosure (resume' tr >>= closureFunc fst')
  
  eofTree :: (Ord a, Eq gamma) => FST a pred func -> PathTree [gamma] a -> PathTree [gamma] a
  eofTree fst' tr = runClosure (resume' tr >>= eof fst')
  
  consumeTree :: (Eq func', Eq var, Eq (Rng func'), PartialOrder pred, Ord st)
              => (func -> func')
              -> FST st pred func
              -> pred
              -> PathTree (UpdateStringFunc var func') st
              -> PathTree (UpdateStringFunc var func') st
  consumeTree inj fst' p tr = runClosure (resume' tr >>= consume inj fst' p)
  
  killTree :: (Ord st, Eq var, Eq func, Eq (Rng func)) =>
              S.Set st -> PathTree (UpdateStringFunc var func) st
                       -> PathTree (UpdateStringFunc var func) st
  killTree kills tr = runClosure (resume' tr >>= kill kills)
  
  
  -------------------------
  -- Abstracting path trees
  -------------------------
  
  abstract :: Tree (UpdateStringFunc Var func) a
              -> ([(Var, UpdateStringFunc Var func)], Tree Var a)
  abstract t = go [] t
    where
    go v (Tip w a) = ([(Var v, w)], Tip (Var v) a)
    go v (Fork w ts) =
      let (ms, ts') = unzip [ go (m:v) t' | (m, t') <- zip [0..] ts ]
      in ((Var v, w):concat ms, Fork (Var v) ts')
  
  abstract' :: PathTree (UpdateStringFunc Var func) a
              -> ([(Var, UpdateStringFunc Var func)], PathTree Var a)
  abstract' Nothing = ([], Nothing)
  abstract' (Just t) = let (m', t') = abstract t in (m', Just t')
  
  unabstract :: PathTree (UpdateString var (Rng func)) a
             -> PathTree (UpdateStringFunc var func) a
  unabstract = fmap (mapOutput (constUpdateStringFunc))
  
  
  ---------------------
  -- SST specialization
  ---------------------
  
  specialize :: (Function func, Dom func ~ [sigma], Rng func ~ [delta]
                ,Enumerable pred sigma) =>
                [(st, [pred], [(var, UpdateStringFunc var func)], st)]
                -> [(st, [pred], [(var, UpdateStringFunc var func)], st)]
  specialize ts =
    [ (q, ps, [ (v, resolveConstFunc $ specFunc ps f) | (v,f) <- xs ], q')
      | (q, ps, xs, q') <- ts ]
    where
      resolveConstFunc = normalizeUpdateStringFunc . concatMap resolve
      resolve (FuncA f) | Just c <- isConst f = [ConstA c]
      resolve x = [x]
  
      specFunc ps f =
        if product (map size ps) == 1 then
            constUpdateStringFunc $ evalUpdateStringFunc (map (lookupIndex 0) ps) f
        else
            f
  
  
  ----------------------------
  -- SST construction from FST
  ----------------------------
  
  -- | Consume a number of abstract symbols (represented as a list of predicates)
  -- and compute the resulting path tree.
  consumeTreeMany :: (PartialOrder pred, Ord st, Eq var, Eq func, Eq (Rng func)) =>
                     FST st pred func
                  -> [pred]
                  -> PathTree (UpdateStringFunc var (SSTFunc func)) st
                  -> PathTree (UpdateStringFunc var (SSTFunc func)) st
  consumeTreeMany fst' preds = go 0 preds
      where
        -- SST functions have list ranges, whereas the underlying FST function is
        -- assumed to look at a single symbol at a time. `inj i f' is a function
        -- that calls f on the i'th symbol in the input list and leaves the rest
        -- alone.
        inj = let n = length preds
              in \i f -> replicate i EpsFunc ++ JustFunc f:replicate (n-i-1) EpsFunc
  
        go _ [] tr = tr
        go i (p:ps) tr = go (i+1) ps (closureTreeFunc fst' (consumeTree (inj i) fst' p (closureTreeFunc fst' tr)))
  
  -- | Determinize the path tree simulation algorithm for some FST
  sstFromFST
    :: (Bounded (Dom func), Enum (Dom func), Eq func,
        Eq (Rng func), Ord st, Ord pred, Monoid (Rng func),
        PartialOrder pred, Function func, Enumerable pred (Dom func),
        Boolean pred, Rng func ~ [delta]) =>
       FST st pred func
       -> Bool -> SST (PathTree Var st) pred (SSTFunc func) Var
  sstFromFST fst' singletonMode =
    construct initialState (specialize transitions) outputs
    where
      initialState           = Just (Tip (Var []) (fstI fst'))
      (transitions, outputs) = saturate (S.singleton initialState) S.empty [] []
  
      saturate ws states trans outs
        | S.null ws                      = (trans, outs)
        | (t, ws') <- S.deleteFindMin ws =
          if S.member t states then saturate ws' states trans outs else
          let tcl  = closureAbstractTree fst' t
              tcl' = unabstract tcl
              os   = [ (t, w) | Tip w _ <- maybeToList $ eofTree fst' tcl ]
              ts   = do (ps, kills) <- prefixTests fst' singletonMode (concatMap tflat $ maybeToList tcl')
                        guard $ not $ null ps
                        let (kappa, t'') = abstract' $ consumeTreeMany fst' ps $ killTree kills tcl'
                        guard (isJust t'')
                        return (t, ps, kappa, t'')
              wl'  = S.fromList [ newt | (_,_,_,newt) <- ts ]
          in saturate (S.union ws' wl') (S.insert t states) (ts ++ trans) (os ++ outs)
  