{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE RankNTypes        #-}

module Precursor.Structure.Traversable
  ( -- * The 'Traversable' class
    Traversable
  , traverse
  , traverse_
  , sequenceA
  , sequence
  , sequence_
    -- * Utility functions
  , for
  , for_
  , mapAccumL
  , mapAccumR
  , zipInto
  ) where

import           Data.Foldable                 (Foldable, for_, sequenceA_,
                                                traverse_)
import           Data.Traversable              hiding (sequence)
import           Precursor.Control.Applicative
import           Precursor.Control.Category
import           Precursor.Control.Functor
import           Precursor.Data.Maybe
import           Precursor.Function
import           Precursor.Structure.Foldable

-- $setup
-- >>> import Test.QuickCheck
-- >>> import Precursor.Numeric.Num

-- | Evaluate each action in the structure from left to right, and
-- and collect the results. For a version that ignores the results
-- see 'sequence_'.
sequence :: (Traversable t, Applicative f) => t (f a) -> f (t a)
sequence = sequenceA

-- | Evaluate each action in the structure from left to right, and
-- ignore the results. For a version that doesn't ignore the results
-- see 'sequence'.
sequence_ :: (Foldable t, Applicative f) => t (f a) -> f ()
sequence_ = sequenceA_

-- | A Scott-encoding of a list. This probably isn't very efficient.
newtype List a =
  List (forall b. b -> (a -> List a -> b) -> b)

newtype State s a =
  State (forall c. (a -> s -> c) -> s -> c)

instance Functor (State s) where
  fmap f (State m) = State (\t -> m (t . f))
  {-# INLINABLE fmap #-}

instance Applicative (State s) where
  pure x = State (\t -> t x)
  {-# INLINABLE pure #-}
  State fs <*> State xs =
    State (\t -> fs (\f -> xs (t . f)))
  {-# INLINABLE (<*>) #-}

evalState :: State s a -> s -> a
evalState (State x) = x const
{-# INLINABLE evalState #-}

-- | Zip two structures together, preserving the shape of the left.
--
-- prop> zipInto const (xs :: [Int]) (ys :: [Int]) === xs
zipInto
  :: (Traversable t, Foldable f)
  => (a -> Maybe b -> c)
  -> t a
  -> f b
  -> t c
zipInto f xs =
  evalState (traverse (flip fmap pop . f) xs) . foldr cons nil where
    cons y ys = List (const (\g -> g y ys))
    nil = List const
    pop = State (\t (List l) -> l (t Nothing nil) (t . Just))
{-# INLINABLE zipInto #-}