55 lines
1.9 KiB
Haskell
55 lines
1.9 KiB
Haskell
{-# LANGUAGE DeriveFunctor #-}
|
|
{-# LANGUAGE DeriveFoldable #-}
|
|
{-# LANGUAGE DeriveTraversable #-}
|
|
{-# LANGUAGE DeriveGeneric #-}
|
|
module Arbitrary where
|
|
|
|
import Test.QuickCheck
|
|
import GHC.Generics (Generic(..))
|
|
import Text.Pretty.Simple
|
|
import Data.List
|
|
|
|
data TreeShape = Root
|
|
{ treeSize :: Int
|
|
, treeSubSizes :: [TreeShape]
|
|
} deriving (Show, Eq)
|
|
|
|
data EnzymeTree s a = EnzymeTree TreeShape a [EnzymeTree TreeShape a] deriving (Show, Eq, Generic, Functor, Foldable, Traversable)
|
|
|
|
instance Applicative (EnzymeTree s) where
|
|
pure a = EnzymeTree (Root 1 []) a []
|
|
EnzymeTree s f fs <*> EnzymeTree _ a as = EnzymeTree s (f a) (zipWith (<*>) fs as)
|
|
|
|
instance Arbitrary a => Arbitrary (EnzymeTree s a) where
|
|
arbitrary = sized arbitrarySizedEnzymeTree
|
|
|
|
arbitrarySizedEnzymeTree :: Arbitrary a => Int -> Gen (EnzymeTree s a)
|
|
arbitrarySizedEnzymeTree m = do
|
|
t <- arbitrary
|
|
n <- choose (m `div` 4, m)
|
|
ts <- if n == 0 then return []
|
|
else sequenceA $ replicate (2*n) (arbitrarySizedEnzymeTree (n `div` 2))
|
|
let (_,ts') = foldr (go m) (0,[]) ts
|
|
go :: Int -> EnzymeTree s a -> (Int, [EnzymeTree s a]) -> (Int, [EnzymeTree s a])
|
|
go m x (s,ts)
|
|
| m == 0 = (s,ts)
|
|
| s + getTreeSize x > m = (s,ts)
|
|
| otherwise = (s + getTreeSize x, x:ts)
|
|
|
|
sz = sum $ getTreeSize <$> ts'
|
|
ss = (\(EnzymeTree a _ _) -> a) <$> ts'
|
|
return (EnzymeTree (Root (sz+1) ss) t ts')
|
|
|
|
getTreeSize :: EnzymeTree s a -> Int
|
|
getTreeSize (EnzymeTree (Root a _) _ _) = a
|
|
|
|
getShape :: EnzymeTree s a -> TreeShape
|
|
getShape (EnzymeTree s _ _) = s
|
|
|
|
treeFromList :: TreeShape -> [a] -> EnzymeTree s a
|
|
treeFromList (Root n ns) (a:as) = EnzymeTree (Root n ns) a (unfoldr go (ns,as))
|
|
where
|
|
go :: ([TreeShape],[a]) -> Maybe (EnzymeTree s a,([TreeShape],[a]))
|
|
go (n:ns,as) = Just (treeFromList n (take (treeSize n) as), (ns, drop (treeSize n) as))
|
|
go ([],_) = Nothing
|