Unfold.hs

{-# LANGUAGE BangPatterns, BlockArguments, GADTs, LambdaCase, ExplicitNamespaces #-}
-- | Implementations of unfolds using weaves.
--
-- Also includes some implementations without weaves for comparison.
--
-- This module is implemented as a Backpack mixin, which imports the signature "Weave.Sig",
-- to be instantiated with "Weave.Strict", "Weave.Lazy", and "Weave.Endless".
module Weave.Unfold where

import Data.Foldable (traverse_)
import Data.Functor ((<&>))
import Data.Functor.Identity
import Data.Tree (Tree(..))
import Data.Fix (Fix(..))
import qualified Data.Tree.Binary.Preorder as Binary
import Weave.Sig

-- | Weave a rose tree.
weave_Tree :: Applicative m => (b -> m (a, [b])) -> b -> m (Weave m (Tree a))
weave_Tree f = go
  where
    go b = f b <&> \(a, bs) ->
      fmap (Node a) (foldr (\b0 bfs -> liftA2 (:) (weft (go b0)) bfs) (pure []) bs)

-- | Unfold a rose tree.
weaveM_BF_Tree :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a)
weaveM_BF_Tree f b = weave_Tree f b >>= mesh

-- | Weave a binary tree.
weave_Binary :: Applicative m => (b -> m (Maybe (a, b, b))) -> b -> m (Weave m (Binary.Tree a))
weave_Binary f = go where
  go b = f b <&> \case
      Nothing -> pure Binary.Leaf
      Just (a, b1, b2) -> (\ ~(x, y) -> Binary.Node a x y) <$> weft ((liftA2 . liftA2) (,) (go b1) (go b2))

-- | Unfold a binary tree.
weaveM_BF_Binary :: Monad m => (b -> m (Maybe (a, b, b))) -> b -> m (Binary.Tree a)
weaveM_BF_Binary f b = weave_Binary f b >>= mesh

-- | Weave an @f@-tree, for any functor @f@.
weaveM_BF_Fix :: (Traversable f, Applicative m) => (b -> m (f b)) -> b -> m (Weave m (Fix f))
weaveM_BF_Fix f b = f b <&> \t -> Fix <$> traverse (weft . weaveM_BF_Fix f) t

-- | Unfold an @f@-tree.
unfoldM_BF_Fix :: (Traversable f, Monad m) => (b -> m (f b)) -> b -> m (Fix f)
unfoldM_BF_Fix f b = weaveM_BF_Fix f b >>= mesh

-- * Pure unfolds

-- | Purely unfold an @f@-tree, for any functor @f@.
--
-- This is included for reference; it doesn't use the 'Weave' abstraction.
--
-- Identical to 'Data.Fix.unfoldFix' from the library @data-fix@.
unfold :: Functor f => (b -> f b) -> b -> Fix f
unfold f b = Fix (unfold f <$> f b)

-- | Depth-first tree unfold.
--
-- This is included for reference; it doesn't use the 'Weave' abstraction.
unfoldM_DF :: (Traversable f, Monad m) => (b -> m (f b)) -> b -> m (Fix f)
unfoldM_DF f b = Fix <$> (f b >>= traverse (unfoldM_DF f))

-- | Same as 'unfold', implemented using 'unfoldM_DF'.
--
-- This is included for reference; it doesn't use the 'Weave' abstraction.
unfold_DF :: Traversable f => (b -> f b) -> b -> Fix f
unfold_DF f = runIdentity . unfoldM_DF (pure . f)

-- | Same as 'unfold', implemented using 'weaveM_BF_Fix'.
unfold_BF :: Traversable f => (b -> f b) -> b -> Fix f
unfold_BF f = runIdentity . (>>= mesh) . weaveM_BF_Fix (pure . f)

-- * Oblivious unfolds

-- | Breadth-first unfold.
unfold_BF_ :: (Foldable f, Applicative m) => (b -> m (f b)) -> b -> m (Weave m ())
unfold_BF_ f b = f b <&> \t -> traverse_ (weft . unfold_BF_ f) t

-- | Depth-first unfold.
--
-- This is included for reference; it doesn't use the 'Weave' abstraction.
unfold_DF_ :: (Foldable f, Monad m) => (b -> m (f b)) -> b -> m ()
unfold_DF_ f b = f b >>= traverse_ (unfold_DF_ f)