{-# LANGUAGE GADTs #-}
-- | Lazy weaves enable linear-time implementations of breadth-first unfolds.
module Weave.Lazy
  ( Weave(..)
  , weft
  , mesh
  ) where

-- | Lazy weaves.
--
-- The 'Applicative' operation @('liftA2')@ combines weaves level-wise.
data Weave m a where
  Pure :: a -> Weave m a
  Weft :: m (Weave m b) -> (b -> a) -> Weave m a

instance Functor (Weave m) where
  fmap :: forall a b. (a -> b) -> Weave m a -> Weave m b
fmap a -> b
f (Pure a
x) = b -> Weave m b
forall a (m :: * -> *). a -> Weave m a
Pure (a -> b
f a
x)
  fmap a -> b
f (Weft m (Weave m b)
u b -> a
g) = m (Weave m b) -> (b -> b) -> Weave m b
forall (m :: * -> *) b a. m (Weave m b) -> (b -> a) -> Weave m a
Weft m (Weave m b)
u (a -> b
f (a -> b) -> (b -> a) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> a
g)

instance Applicative m => Applicative (Weave m) where
  pure :: forall a. a -> Weave m a
pure = a -> Weave m a
forall a (m :: * -> *). a -> Weave m a
Pure
  liftA2 :: forall a b c. (a -> b -> c) -> Weave m a -> Weave m b -> Weave m c
liftA2 a -> b -> c
f (Pure a
x) Weave m b
u = (b -> c) -> Weave m b -> Weave m c
forall a b. (a -> b) -> Weave m a -> Weave m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> b -> c
f a
x) Weave m b
u
  liftA2 a -> b -> c
f (Weft m (Weave m b)
u b -> a
g) (Pure b
y) = m (Weave m b) -> (b -> c) -> Weave m c
forall (m :: * -> *) b a. m (Weave m b) -> (b -> a) -> Weave m a
Weft m (Weave m b)
u (\b
x -> a -> b -> c
f (b -> a
g b
x) b
y)
  liftA2 a -> b -> c
f (Weft m (Weave m b)
u b -> a
g) (Weft m (Weave m b)
v b -> b
h) = m (Weave m (b, b)) -> ((b, b) -> c) -> Weave m c
forall (m :: * -> *) b a. m (Weave m b) -> (b -> a) -> Weave m a
Weft ((Weave m b -> Weave m b -> Weave m (b, b))
-> m (Weave m b) -> m (Weave m b) -> m (Weave m (b, b))
forall a b c. (a -> b -> c) -> m a -> m b -> m c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 ((b -> b -> (b, b)) -> Weave m b -> Weave m b -> Weave m (b, b)
forall a b c. (a -> b -> c) -> Weave m a -> Weave m b -> Weave m c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,)) m (Weave m b)
u m (Weave m b)
v) (\ ~(b
x, b
y) -> a -> b -> c
f (b -> a
g b
x) (b -> b
h b
y))

-- | A weft is one level of 'Weave'. It is a computation which returns the remaining levels.
weft :: m (Weave m a) -> Weave m a
weft :: forall (m :: * -> *) a. m (Weave m a) -> Weave m a
weft m (Weave m a)
u = m (Weave m a) -> (a -> a) -> Weave m a
forall (m :: * -> *) b a. m (Weave m b) -> (b -> a) -> Weave m a
Weft m (Weave m a)
u a -> a
forall a. a -> a
id

-- | Run all the wefts in a 'Weave' sequentially.
mesh :: Monad m => Weave m a -> m a
mesh :: forall (m :: * -> *) a. Monad m => Weave m a -> m a
mesh (Pure a
x) = a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x
mesh (Weft m (Weave m b)
u b -> a
f) = b -> a
f (b -> a) -> m b -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (m (Weave m b)
u m (Weave m b) -> (Weave m b -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Weave m b -> m b
forall (m :: * -> *) a. Monad m => Weave m a -> m a
mesh)