1 files changed, 187 insertions, 0 deletions

diff --git a/edit-lens/test/Control/FSTTest.hs b/edit-lens/test/Control/FSTTest.hs
new file mode 100644
index 0000000..f5e02c2
--- /dev/null
+++ b/edit-lens/test/Control/FSTTest.hs
@@ -0,0 +1,187 @@
+module Control.FSTTest where
+
+import Control.FST
+
+import Data.Set (Set)
+import qualified Data.Set as Set
+
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
+
+import Data.Sequence (Seq)
+import qualified Data.Sequence as Seq
+
+import Data.Maybe (fromMaybe)
+import Data.Foldable (Foldable(..))
+
+import Control.Monad (when)
+
+import Data.Void
+
+import Test.Tasty
+import Test.Tasty.Hedgehog
+import Test.Tasty.HUnit hiding (assert)
+
+import Hedgehog
+import qualified Hedgehog.Gen as G
+import qualified Hedgehog.Range as R
+
+import Numeric.Natural
+
+import Text.PrettyPrint.Leijen (Pretty(..))
+
+import Control.DeepSeq (force)
+
+
+fstId :: Ord a => Set a -> FST () a a
+fstId syms = FST
+  { stInitial = Set.singleton ()
+  , stTransition = Map.fromList [(((), Just sym), Set.singleton ((), Just sym)) | sym <- Set.toList syms]
+  , stAccept = Set.singleton ()
+  }
+
+fstDouble :: Ord a => Set a -> FST (Maybe a) a a
+fstDouble syms = FST
+  { stInitial = Set.singleton Nothing
+  , stTransition = Map.fromListWith Set.union . concat $
+    [ [((Nothing, Just sym), Set.singleton (Just sym, Just sym)) | sym <- Set.toList syms]
+    , [((Just sym, Nothing), Set.singleton (Nothing, Just sym)) | sym <- Set.toList syms]
+    ]
+  , stAccept = Set.singleton Nothing
+  }
+
+fstRunLengthDecode :: Ord a
+                   => Set a -- ^ Alphabet
+                   -> Natural -- ^ Upper limit to run length
+                   -> FST (Maybe (Natural, Maybe a)) (Either Natural a) a
+fstRunLengthDecode syms lim = FST
+  { stInitial = Set.singleton Nothing
+  , stTransition = Map.fromListWith Set.union . concat $
+    [ [((Nothing, Just (Left n)), Set.singleton (Just (n, Nothing), Nothing)) | n <- [0..lim]]
+    , [((Just (n, Nothing), Just (Right sym)), Set.singleton (Just (n, Just sym), Nothing)) | n <- [0..lim], sym <- Set.toList syms]
+    , [((Just (n, Just sym), Nothing), Set.singleton (Just (pred n, Just sym), Just sym)) | n <- [1..lim], sym <- Set.toList syms]
+    , [((Just (0, Just sym), Nothing), Set.singleton (Nothing, Nothing)) | sym <- Set.toList syms]
+    ]
+  , stAccept = Set.singleton Nothing
+  }
+
+data StRunLengthEncode a = STREInitial
+                         | STRECountUp a Natural
+                         | STRESwitch (Maybe a) a
+                         | STREFinish
+  deriving (Show, Eq, Ord)
+
+fstRunLengthEncode :: Ord a
+                   => Bool -- ^ Generate /all/ run length encodings instead of the best
+                   -> Set a -- ^ Alphabet
+                   -> Natural -- ^ Upper limit to run length
+                   -> FST (StRunLengthEncode a) a (Either Natural a)
+fstRunLengthEncode genAll syms lim = FST
+  { stInitial = Set.singleton STREInitial
+  , stTransition = Map.fromListWith Set.union . concat $
+    [ [((STREInitial, Just sym), Set.singleton (STRECountUp sym 1, Nothing)) | sym <- Set.toList syms]
+    , [((STRECountUp sym n, Just sym), Set.singleton (STRECountUp sym (succ n), Nothing)) | sym <- Set.toList syms, n <- [1..pred lim]]
+    , [((STRECountUp sym n, Just sym'), Set.singleton (STRESwitch (Just sym') sym, Just $ Left n)) | n <- [1..lim], sym <- Set.toList syms, sym' <- Set.toList syms, sym /= sym']
+    , [((STRECountUp sym lim, Just sym), Set.singleton (STRESwitch (Just sym) sym, Just $ Left lim)) | sym <- Set.toList syms]
+    , [((STRECountUp sym n, Nothing), Set.singleton (STRESwitch Nothing sym, Just $ Left n)) | sym <- Set.toList syms, n <- [1..lim]]
+    , [((STRESwitch (Just sym') sym, Nothing), Set.singleton (STRECountUp sym' 1, Just $ Right sym)) | sym <- Set.toList syms, sym' <- Set.toList syms]
+    , [((STRESwitch Nothing sym, Nothing), Set.singleton (STREFinish, Just $ Right sym)) | sym <- Set.toList syms]
+    , [((STRECountUp sym n, Just sym), Set.singleton (STRESwitch (Just sym) sym, Just $ Left n)) | n <- [1..lim], sym <- Set.toList syms, genAll]
+    ]
+  , stAccept = Set.fromList [STREInitial, STREFinish]
+  }
+
+
+genWord :: Gen [Natural]
+genWord = G.list (R.linear 0 1000) . G.integral $ R.linear 0 100
+
+  
+runFSTDet :: (Show output, Ord output, Show state, Ord state) => (Set Natural -> FST state Natural output) -> (Seq Natural -> Seq output) -> Property
+runFSTDet mkFst f = property $ do
+  input <- forAll genWord
+  let fst = mkFst $ Set.fromList input
+  annotateShow $ pretty fst
+  [f $ Seq.fromList input] === runFST fst (Seq.fromList input)
+
+
+hprop_runFSTId, hprop_runFSTDouble :: Property
+hprop_runFSTId = runFSTDet fstId id
+hprop_runFSTDouble = runFSTDet fstDouble double
+  where
+    double :: Seq a -> Seq a
+    double Seq.Empty = Seq.Empty
+    double (a Seq.:<| as) = a Seq.:<| a Seq.:<| double as
+
+hprop_runWordFST :: Property
+hprop_runWordFST = property $ do
+  input <- forAll genWord
+  let fst = wordFST $ Seq.fromList input
+  annotateShow $ pretty fst
+  [Seq.fromList input] === runFST fst (Seq.empty :: Seq Void)
+
+unit_runLengthDecode, unit_runLengthEncode :: Assertion
+unit_runLengthDecode = runFST fst input @?= [Seq.fromList "aaacc"]
+  where
+    input = Seq.fromList [Left 3, Right 'a', Left 0, Right 'b', Left 2, Right 'c']
+    fst = fstRunLengthDecode (Set.fromList "abc") 3
+unit_runLengthEncode = runFST fst input @?= [Seq.fromList [Left 3, Right 'a', Left 2, Right 'c']]
+  where
+    input = Seq.fromList "aaacc"
+    fst = fstRunLengthEncode False (Set.fromList "abc") 3
+
+hprop_runLength :: Property
+hprop_runLength = property $ do
+  input <- forAll genWord
+  let maxRun = fromMaybe 1 $ countMaxRun input
+      alphabet = Set.fromList input
+      encode = fstRunLengthEncode False alphabet maxRun
+      decode = fstRunLengthDecode alphabet maxRun
+  annotateShow $ countMaxRun input
+  annotateShow $ pretty encode
+  annotateShow $ pretty decode
+  [encoded] <- return . runFST encode $ Seq.fromList input
+  annotateShow $ toList encoded
+  [Seq.fromList input] === runFST decode encoded
+
+hprop_runLength' :: Property
+hprop_runLength' = property $ do
+  input <- forAll genWord
+
+  let maxRun = fromMaybe 1 $ countMaxRun input
+      alphabet = Set.fromList input
+      encode = fstRunLengthEncode True alphabet maxRun
+      decode = fstRunLengthDecode alphabet maxRun
+  annotateShow $ countMaxRun input
+  annotateShow $ pretty encode
+  annotateShow $ pretty decode
+
+  let encoded = runFST encode $ Seq.fromList input
+  annotateShow encoded
+  when (maxRun > 1) $
+    assert $ encoded `longerThan` 1
+
+  encoded' <- forAll $ G.element encoded
+  [Seq.fromList input] === runFST decode encoded'
+  where
+    longerThan :: [a] -> Natural -> Bool
+    longerThan []     _ = False
+    longerThan  _     0 = True
+    longerThan (_:xs) n = longerThan xs $ pred n
+
+countMaxRun :: (Eq a, Foldable f) => f a -> Maybe Natural
+countMaxRun = goMaxRun Nothing Nothing . toList
+  where
+    goMaxRun :: Eq a
+             => Maybe (a, Natural) -- ^ Maximum
+             -> Maybe (a, Natural) -- ^ Current
+             -> [a]
+             -> Maybe Natural
+    goMaxRun Nothing Nothing [] = Nothing
+    goMaxRun (Just (_, n)) (Just (_, n')) []
+      | n' > n    = Just n'
+      | otherwise = Just n
+    goMaxRun Nothing Nothing (a:as) = goMaxRun (Just (a, 1)) (Just (a, 1)) as
+    goMaxRun (Just (a, n)) (Just (a', n')) (x:xs)
+      | x == a'   = goMaxRun (Just (a, n)) (Just (a', succ n')) xs
+      | n' > n    = goMaxRun (Just (a', n')) (Just (x, 1)) xs
+      | otherwise = goMaxRun (Just (a, n)) (Just (x, 1)) xs