{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}

-- |
-- Module      : Data.Diagram
-- Copyright   : (c) Eddie Jones 2020
-- License     : BSD-3
-- Maintainer  : eddiejones2108@gmail.com
-- Stability   : experimental
--
-- This module contains a 'Free' /monad/ with shared subterms.
-- Unfortunately it is not a true monad in /Hask/ category.
-- If the classical Free monad generalises a tree, here we generalise directed acylic graphs.
-- The 'Diagram' is the true monadic context in which all operations are run.
--
-- For most operations to work the functor must be of class 'Eq1', 'Hashable1' and 'Traversable'.
module Data.Diagram
  ( -- * Diagrams
    Diagram,
    runDiagram,
    compress,

    -- * Nodes
    Free (Pure),
    free,
    fromFree,
    retract,

    -- * Combinators
    map,
    bind,
    fold,
  )
where

import Control.Monad.State hiding (foldM)
import Data.Functor.Classes
import Data.Functor.Identity
import qualified Data.HashMap.Lazy as H
import Data.Hashable
import Data.Hashable.Lifted
import qualified Data.IntMap as I
import GHC.Generics (Generic)
import Prelude hiding (lookup, map)

-- | The diagram records overlapping terms:
--
--      * @f@ is the functor from which the Free terms are build.
--
--      * @a@ is the type of leaves or terminals.
--
--      * @s@ is a phantom type labelling the diagram and preventing leaking.
newtype Diagram f a s m r = Diagram
  { getState ::
      StateT
        ( Int, -- Fresh node ID
          H.HashMap (Free_ f a s) Int, -- node to ID
          I.IntMap (Free_ f a s) -- ID to node
        )
        m
        r
  }
  deriving (Functor, Applicative, Monad, MonadTrans)

-- | Extract non-diagrammatic information
runDiagram :: Monad m => (forall s. Diagram f a s m b) -> m b
runDiagram (Diagram d) = evalStateT d (0, H.empty, I.empty)

-- | Remove diagrmmatic information from the underlying monad.
-- This could have problematic side-effects and should be used with caution
compress :: Monad m => Diagram f a s m b -> Diagram f a s Identity (m b)
compress = Diagram . gets . evalStateT . getState

-- Find a node from it's ID
lookup :: Monad m => Int -> Diagram f a s m (f (Free f a s))
lookup i = Diagram $ do
  (_, _, fs) <- get
  case I.lookup i fs of
    Nothing -> error "No free with that id!"
    Just f -> return (inner f)

-- | The Free monad parameterised by it's diagram.
-- Approximately it corresponds to a non-terminal or sub-diagram
data Free (f :: * -> *) a s
  = Pure a
  | ID Int
  deriving (Eq, Ord, Generic)

newtype Free_ f a s = Free
  { inner :: f (Free f a s)
  }

instance Hashable a => Hashable (Free f a s)

instance (Eq a, Eq1 f) => Eq (Free_ f a s) where
  Free f == Free g = liftEq (==) f g

instance (Hashable a, Hashable1 f) => Hashable (Free_ f a s) where
  hashWithSalt s (Free f) = liftHashWithSalt hashWithSalt s f

-- | Analogous to the Free constructor, this fuction wraps another functorial layer.
-- In a graph setting, this corresponds to creating a node from it's offspring structure
free :: (Monad m, Hashable1 f, Hashable a, Eq1 f, Eq a) => f (Free f a s) -> Diagram f a s m (Free f a s)
free f = Diagram $ do
  (i, h, fs) <- get
  case H.lookup (Free f) h of
    Just j -> return (ID j) -- Node sharing
    Nothing -> do
      put (i + 1, H.insert (Free f) i h, I.insert i (Free f) fs)
      return (ID i)

-- | Case analysis for 'Free', the first argument is the 'free' continuation, and the second is the 'Pure' case.
{-# INLINE fromFree #-}
fromFree :: Monad m => Free f a s -> (f (Free f a s) -> Diagram f a s m b) -> (a -> Diagram f a s m b) -> Diagram f a s m b
fromFree (Pure a) _ b = b a
fromFree (ID i) f _ = lookup i >>= f

-- | When @f@ is monadic itself, 'retract' will collapse the recursive structure
retract :: (Monad f, Traversable f) => Free f a s -> Diagram f a s f a
{-# INLINE retract #-}
retract = fold lift (lift . return)

-- | Apply a function to all terminals in a 'Free' sub-diagram
{-# INLINE map #-}
map :: (Hashable1 f, Hashable a, Eq1 f, Eq a, Monad m, Traversable f) => (a -> Diagram f a s m a) -> Free f a s -> Diagram f a s m (Free f a s)
map f = fold free (fmap Pure . f)

-- | Replace a terminal with a non-terminal in a 'Free' sub-diagram
{-# INLINE bind #-}
bind :: (Hashable1 f, Hashable a, Eq1 f, Eq a, Monad m, Traversable f) => Free f a s -> (a -> Diagram f a s m (Free f a s)) -> Diagram f a s m (Free f a s)
bind = flip (fold free)

-- | Collapses a 'Free' sub-diagram by instianting the terminals and offpsring structure
{-# INLINE fold #-}
fold :: (Monad m, Traversable f) => (f b -> Diagram f a s m b) -> (a -> Diagram f a s m b) -> Free f a s -> Diagram f a s m b
fold f b x = fromFree x (mapM (fold f b) >=> f) b