{-# LANGUAGE RankNTypes #-} module Test.RadixNTree.Word8.Sample ( Sample , zero , one , tip , bin , tiny , small , medium --, large , mkUnary0 , mkUnary1 , mkUnary2 , mkBinary , mkBinaryL , mkEqual , mkSuperset , mkSubset ) where import No.Tree (NoTree) import qualified No.Tree as No import Test.Kit import Test.Random import Data.List.NonEmpty (NonEmpty (..)) import qualified Data.List.NonEmpty as NonEmpty import Data.Word import System.Random data Trees = Trees (NonEmpty Tree) | End deriving Show data Tree = Tree (NonEmpty Word8) Bool -- ^ Whether this point is a separate key in the tree Trees deriving Show genTrees :: RandomGen g => Int -- ^ Maximum branches on each level -> Int -- ^ Maximum number of segments -> Int -- ^ Maximum segment length -> Int -- ^ Maximum total length -> g -> (Trees, g) genTrees nB nL nS nT = broad nL nT where broad count len g0 | len <= 0 || count <= 0 = (End, g0) | otherwise = let ~(n, g1) = uniformR (1, nB) g0 ~(as, g2) = list1 (deep count len) n g1 in (Trees $ dedup as, g2) dedup = NonEmpty.nubBy (\(Tree (x :| _) _ _) (Tree (y :| _) _ _) -> x == y) deep count len g0 = let ~(n, g1) = uniformR (1, max 1 (min len nS)) g0 ~(xs, g2) = list1 uniform n g1 ~(t, g3) = broad (count - 1) (len - n) g2 ~(bias, g4) = case t of End -> (nL, g3) _ -> uniformR (1, nL) g3 in (Tree xs (bias == 1) t, g4) timber :: Trees -> [([Word8], Int)] timber = fst . broad id ([], 1) where broad pre ~(acc, n) End = ((pre [], n) : acc, n + 1) broad pre z (Trees ts) = foldr (flip $ deep pre) z ts deep pre z@(acc, n) (Tree xs real t) = let z' = if real then ((pre $ NonEmpty.toList xs, n) : acc, n + 1) else z in broad (pre . (NonEmpty.toList xs <>)) z' t data Sample = Sample [(No.Openness, [Word8], Int)] -- ^ Keys in the dictionary [(No.Openness, [Word8], Int)] -- ^ Keys not in the dictionary deriving Show halve :: [a] -> ([a], [a]) halve (a:b:cs) = let ~(xs, ys) = halve cs in (a:xs, b:ys) halve a = (a, []) sample :: RandomGen g => Trees -> g -> Sample sample t g0 = let ~(xs, g1) = shuffle (timber t) g0 ~(os, g2) = list (\g' -> let ~(b, g'') = uniform g' in ( if b then No.Open else No.Closed , g'' ) ) (length xs) g1 xs' = zipWith (\a (b, c) -> (a, b, c)) os xs ~(ys, zs) = halve xs' ~(z, _) = uniform g2 as | z = case ys of [] -> [] (b, _, i) : rest -> (b, [], i) : rest | otherwise = ys bs = case zs of [] -> [] (b, _, i) : rest -> (b, [], i) : rest in Sample as bs zero, one, tip, bin :: Sample zero = Sample [] [ (No.Open, [], 1), (No.Closed, [1, 2, 3], 2), (No.Open, [3, 2, 1], 3) ] one = Sample [(No.Open, [1, 2, 3], 0)] [ (No.Closed, [1, 2, 3], 1), (No.Open, [1, 2, 2], 2), (No.Closed, [1, 2, 4], 3) , (No.Open, [1, 2], 4), (No.Closed, [1, 2, 3, 4], 5), (No.Open, [2, 3, 4], 6) , (No.Closed, [], 7), (No.Open, [2], 8) ] tip = Sample [(No.Open, [], 0)] [ (No.Closed, [1, 2, 3], 1), (No.Closed, [], 2) ] bin = Sample [(No.Open, [1, 2, 2, 3], 0), (No.Closed, [1, 2, 4, 5], 1)] [ (No.Closed, [1, 2, 3, 4], 2), (No.Open, [1, 2, 2, 3], 3) , (No.Closed, [1, 2, 4, 5], 4), (No.Closed, [], 5) ] tiny, small, medium :: Sample tiny = uncurry sample $ genTrees 4 2 4 16 (mkStdGen 2) small = uncurry sample $ genTrees 4 4 4 16 (mkStdGen 4) medium = uncurry sample $ genTrees 8 4 4 16 (mkStdGen 16) type FromList pat = forall x. [([Word8], x)] -> pat x mkUnary0 :: FromList pat -> Sample -> [Case () (pat Int) (NoTree [Word8] Int)] mkUnary0 patFromList (Sample xs _) = let as = fmap (\(_, k, i) -> (k, i)) xs in [Case () (patFromList as) (No.fromList as)] mkUnary1 :: FromList pat -> Sample -> [Case (No.Openness, [Word8], Int) (pat Int) (NoTree [Word8] Int)] mkUnary1 patFromList (Sample xs ys) = let as = fmap (\(_, k, i) -> (k, i)) xs pat = patFromList as no = No.fromList as in foldr (\x -> (:) (Case x pat no)) [] $ xs <> ys mkUnary2 :: FromList pat -> Sample -> [Case (No.Openness, [Word8], No.Openness, [Word8], Int) (pat Int) (NoTree [Word8] Int)] mkUnary2 patFromList (Sample xs ys) = let xs' = fmap (\(_, k, i) -> (k, i)) xs pat = patFromList xs' no = No.fromList xs' ~(as, bs) = halve xs ~(cs, ds) = halve ys ones = fmap (\(o, a, i) -> (o, a, o, a, i)) $ as <> cs twos = zipWith (\(o, a, i) (p, b, _) -> (o, a, p, b, i)) bs ds in foldr (\x -> (:) (Case x pat no)) [] $ ones <> twos mkBinary :: FromList pat -> Sample -> [Case (pat Int, NoTree [Word8] Int) (pat Int) (NoTree [Word8] Int)] mkBinary patFromList (Sample xs ys) = let as = fmap (\(_, k, i) -> (k, i)) xs bs = fmap (\(_, k, i) -> (k, i)) ys in [Case (patFromList bs, No.fromList bs) (patFromList as) (No.fromList as)] mkBinaryL :: FromList pat -> Sample -> [Case (pat Int, NoTree [Word8] Int) (pat Int) (NoTree [Word8] Int)] mkBinaryL patFromList (Sample xs ys) = let xs' = fmap (\(_, k, i) -> (k, i)) xs ys' = fmap (\(_, k, i) -> (k, i)) ys ~(as, _) = halve xs' ~(bs, _) = halve ys' ls = fmap (\(k, a) -> (k, negate a)) bs <> xs' rs = fmap (\(k, a) -> (k, negate a)) as <> ys' in [Case (patFromList rs, No.fromList rs) (patFromList ls) (No.fromList ls)] mkEqual :: FromList pat -> Sample -> [Case (pat Int, NoTree [Word8] Int) (pat Int) (NoTree [Word8] Int)] mkEqual patFromList (Sample xs _) = let as = fmap (\(_, k, i) -> (k, i)) xs pat = patFromList as no = No.fromList as in [Case (pat, no) pat no] mkSuperset :: FromList pat -> Sample -> [Case (pat Int, NoTree [Word8] Int) (pat Int) (NoTree [Word8] Int)] mkSuperset patFromList (Sample xs ys) = let as = fmap (\(_, k, i) -> (k, i)) xs bs = fmap (\(_, k, i) -> (k, i)) ys zs = as <> bs in [Case (patFromList zs, No.fromList zs) (patFromList as) (No.fromList as)] mkSubset :: FromList pat -> Sample -> [Case (pat Int, NoTree [Word8] Int) (pat Int) (NoTree [Word8] Int)] mkSubset patFromList (Sample xs ys) = let as = fmap (\(_, k, i) -> (k, i)) xs bs = fmap (\(_, k, i) -> (k, i)) ys zs = as <> bs in [Case (patFromList as, No.fromList as) (patFromList zs) (No.fromList zs)]