added EnzymeTree-Generator via QuickCheck

src/Arbitrary.hs

@@ -0,0 +1,54 @@
+{-# 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