Skip to content

Commit e12fda7

Browse files
committed
Use level-wise algorithm for BF unfolds
`unfoldTreeM_BF` and `unfoldForestM_BF` previously used a queue-based algorithm from a note by Okasaki. That same note indicates that the level-wise approach tends to be slightly faster, and my own informal tests suggest the same. Both solutions suffer from a potential space leak, retaining the seed lists longer than necessary in order to later use their lengths. It remains to be seen whether this leak can be plugged without a speed penalty.
1 parent d195ff2 commit e12fda7

File tree

1 file changed

+24
-36
lines changed

1 file changed

+24
-36
lines changed

Data/Tree.hs

+24-36
Original file line numberDiff line numberDiff line change
@@ -38,14 +38,12 @@ module Data.Tree(
3838
import Data.Foldable (toList)
3939
#else
4040
import Control.Applicative (Applicative(..), (<$>))
41-
import Data.Foldable (Foldable(foldMap), toList)
41+
import Data.Foldable (Foldable(foldMap))
4242
import Data.Monoid (Monoid(..))
4343
import Data.Traversable (Traversable(traverse))
4444
#endif
4545

4646
import Control.Monad (liftM)
47-
import Data.Sequence (Seq, empty, singleton, (<|), (|>), fromList,
48-
ViewL(..), ViewR(..), viewl, viewr)
4947
import Data.Typeable
5048
import Control.DeepSeq (NFData(rnf))
5149

@@ -163,37 +161,27 @@ unfoldForestM :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a)
163161
#endif
164162
unfoldForestM f = Prelude.mapM (unfoldTreeM f)
165163

166-
-- | Monadic tree builder, in breadth-first order,
167-
-- using an algorithm adapted from
168-
-- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/,
169-
-- by Chris Okasaki, /ICFP'00/.
164+
-- | Monadic tree builder, in breadth-first order.
170165
unfoldTreeM_BF :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a)
171-
unfoldTreeM_BF f b = liftM getElement $ unfoldForestQ f (singleton b)
172-
where
173-
getElement xs = case viewl xs of
174-
x :< _ -> x
175-
EmptyL -> error "unfoldTreeM_BF"
176-
177-
-- | Monadic forest builder, in breadth-first order,
178-
-- using an algorithm adapted from
179-
-- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/,
180-
-- by Chris Okasaki, /ICFP'00/.
181-
unfoldForestM_BF :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a)
182-
unfoldForestM_BF f = liftM toList . unfoldForestQ f . fromList
183-
184-
-- takes a sequence (queue) of seeds
185-
-- produces a sequence (reversed queue) of trees of the same length
186-
unfoldForestQ :: Monad m => (b -> m (a, [b])) -> Seq b -> m (Seq (Tree a))
187-
unfoldForestQ f aQ = case viewl aQ of
188-
EmptyL -> return empty
189-
a :< aQ' -> do
190-
(b, as) <- f a
191-
tQ <- unfoldForestQ f (Prelude.foldl (|>) aQ' as)
192-
let (tQ', ts) = splitOnto [] as tQ
193-
return (Node b ts <| tQ')
194-
where
195-
splitOnto :: [a'] -> [b'] -> Seq a' -> (Seq a', [a'])
196-
splitOnto as [] q = (q, as)
197-
splitOnto as (_:bs) q = case viewr q of
198-
q' :> a -> splitOnto (a:as) bs q'
199-
EmptyR -> error "unfoldForestQ"
166+
unfoldTreeM_BF f b0 = do
167+
(a, bs) <- f b0
168+
Node a `liftM` unfoldForestM_BF f bs
169+
170+
-- | Monadic forest builder, in breadth-first order.
171+
unfoldForestM_BF :: Monad m
172+
=> (b -> m (a, [b])) -> [b] -> m (Forest a)
173+
unfoldForestM_BF _f [] = return []
174+
unfoldForestM_BF f bs = do
175+
asbss' <- mapM f bs
176+
rebuild asbss' `liftM` unfoldForestM_BF f (concatMap snd asbss')
177+
where
178+
rebuild :: [(a, [any])] -> [Tree a] -> [Tree a]
179+
rebuild [] ts = ts
180+
rebuild ((a, bs') : xs) ts =
181+
case splitAtLength bs' ts of
182+
(us, ts') -> Node a us : rebuild xs ts'
183+
184+
splitAtLength :: [any] -> [a] -> ([a],[a])
185+
splitAtLength (_ : n) (x : xs) = (x : early, late)
186+
where (early, late) = splitAtLength n xs
187+
splitAtLength _ xs = ([], xs)

0 commit comments

Comments
 (0)