Cached evaluation

DAG.hs

{-# OPTIONS_GHC -Wno-name-shadowing #-}

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
module Data.DAG where

import Data.Kind (Type)
import Prelude hiding ((!!))
import qualified Fcf as FCF

-- >>> evaluate test
-- Hey!
-- Hey!
-- Hey!
-- You!
-- Hey!
-- Tom!
--
-- >>> evaluate' test
-- Hey!
-- You!
-- Tom!
test :: DagF IO (() ':> () ':> () ':> 'VNil)
test
  = Node (SFZ :& SFS SFZ :& HNil) (\_ -> putStrLn "Tom!")
  $ Node (SFZ :& SFZ :& SFZ :& HNil) (\_ -> putStrLn "You!")
  $ Node HNil (\_ -> putStrLn "Hey!")
  $ Empty

-------------------------------------------------------------------------------

-- | We need a unary representation of @Nat@ (as opposed to, say, the pretty
-- one in GHC.TypeLits) for the sake of the '(!!)' and '(!!!)' families: the
-- pretty representation would mean, for example, that the type of @FZ@ is @Fin
-- (1 + limit)@, which involves a type family (@(+)@). This leads to an illegal
-- type family application in the instance head of '(!!)' and '(!!!)'.
data Nat = Z | S Nat

-- | Unary representation of numbers up-to-and-excluding a given limit. Zero is
-- smaller than the successor of any number (@FZ :: Fin ('S n)@), and the
-- successor of a number smaller than @x@ is a number smaller than @S x@.
data Fin (limit :: Nat) where
  FZ ::              Fin ('S limit)
  FS :: Fin limit -> Fin ('S limit)

-- | With some magic unlocked by @TypeInType@, we can promote the above GADT
-- and use it as an index, allowing us to create a singleton for the 'Fin'
-- type! If your brain hasn't melted: @SFZ :: SFin 'FZ :: Fin ('S n)@.
data SFin (n :: Fin limit) where
  SFZ ::           SFin  'FZ
  SFS :: SFin n -> SFin ('FS n)

-- | A length-indexed vector is a list indexed by the number of inhabitants.
-- One can provide a safe analogue to @(!!)@ by using the 'Fin' type described
-- above: @(!!) :: Vector a n -> Fin n -> a@. As the index is guaranteed to be
-- within bounds, there's no 'Maybe'!
data Vector (a :: Type) (n :: Nat) where
  VNil ::                    Vector a  'Z
  (:>) :: a -> Vector a n -> Vector a ('S n)

infixr :>

-- | A heterogeneous list is one in which the types of its inhabitants needn't
-- all be the same. To this end, we index an @HList@ by the list of its
-- inhabitants. There are a couple particularly weird things happening here,
-- though:
--
-- - We index by a @Vector k n@ rather than a @[k]@. The reason is that, unlike
--   the @(!!)@ described above in the 'Vector' introduction, we can't write
--   this function for @HList@s because we need to know /exactly/ which index
--   is being used to know the result type! This, ta-da, is why we need the
--   singleton for 'Fin', and using a 'Vector' as our index means we can
--   effectively write the @(!!)@ function described above at the type-level
--   with a type family.
--
-- - The @f@ parameter has a weird kind. What we're doing is exploiting some of
--   the tricks outlined in Lysxia's blog[1]: we are indexing our @HList_@ by a
--   constructor that we use as a function to figure out the type of each
--   argument to our @HList@. If this is confusing, look at the type synonyms
--   below.
--
--   [1] https://blog.poisson.chat/posts/2018-08-06-one-type-family.html
data HList_ (f :: k -> FCF.Exp Type) (xs :: Vector k n) where
  HNil ::                                  HList_ f  'VNil
  (:&) :: FCF.Eval (f x) -> HList_ f xs -> HList_ f (x ':> xs)

infixr :&

-- | In the case of a simple @HList@, the values within aren't wrapped in any
-- special type: @HList '[Int, String, Bool]@ is just a three-element list
-- containing an @Int@, a @String@, and a @Bool@! Luckily, the 'Fcf.Pure'
-- function is just what we need: it takes whatever it's given, and returns it
-- with an @Eval@-friendly wrapper. In other words, @Eval (Pure x)@ is... @x@!
-- This means that any function involving an @HList@ can be totally ignorant of
-- the more general 'HList_' structure.
type HList = HList_ FCF.Pure

-- | Sometimes, we have an 'HList' in which all the values are wrapped in some
-- common functor. Imagine, for example, @HList '[IO Bool, IO String, IO ()]@.
-- The 'HListF' type allows us to write this as @HList IO '[Bool, String, IO]@,
-- thanks to the 'Fcf.Pure1' function: @Eval (Pure1 f x)@ gives us @f x@!
type HListF f = HList_ (FCF.Pure1 f)

-- | As a little example, we can "extract" the functor from an 'HListF' using a
-- makeshift 'sequence' operation: we use the @Applicative f@ to lift the
-- '(:&)' operations over @f@.
hsequence :: Applicative f => HListF f xs -> f (HList xs)
hsequence = \case
  HNil    -> pure HNil
  x :& xs -> (:&) <$> x <*> hsequence xs

-- | Just like the @(!!)@ signature we saw with 'Vector', we can promote that
-- exact function to the type level thanks to Lysxia's tricks in a way that we
-- can partially apply, meaning... well, weirder and weirder HLists!
data (!!) :: Vector k n -> Fin n -> FCF.Exp k
type instance FCF.Eval ((x ':> xs) !! 'FZ  ) = x
type instance FCF.Eval ((x ':> xs) !! 'FS n) = FCF.Eval (xs !! n)

-- | For the sake of some function signatures, defining a fully-saturated
-- synonym means we don't end up with quite so many mentions of 'Eval' all over
-- the place.
type xs !!. index = FCF.Eval (xs !! index)

-- | We define the DAG structure similarly to the 'HList_', with an
-- element-transforming function and a list of the nodes available. Here's
-- where we see why we bothered with the 'SFin', though! A node is a list of
-- indices into the tail, a function /from/ the indexed nodes to the "head
-- type", and the tail itself. 
data Dag_ (f :: k -> FCF.Exp Type) (elements :: Vector k n) where
  Empty
    :: Dag_ f 'VNil

  Node
    :: HListF SFin indices
    -> (HList_ ((!!) xs) indices -> FCF.Eval (f x))
    -> Dag_ f xs
    -> Dag_ f (x ':> xs)

-- | A pure DAG (i.e. one without the @f@ context) is defined just like a pure
-- 'HList': with the help of 'Fcf.Pure'!
type Dag = Dag_  FCF.Pure

-- | Absolutely no prizes for guessing how we define a functor-indexed DAG...
type DagF f = Dag_ (FCF.Pure1 f)

-- | We can look up an arbitrary node in the DAG if we're happy to compute the
-- dependency tree. Note that this will recompute dependencies: if I have a
-- node in my DAG that is a transitive dependency 5 times within the sub-graph
-- for my selected node, it will be computed 5 times! See below for a memoised
-- version of this function.
(!!) :: Monad f => DagF f xs -> SFin idx -> f (xs !!. idx)
(!!) node  SFZ    = case node of Node indices f tail -> tail !=! indices >>= f
(!!) node (SFS n) = case node of Node _       _ tail -> tail !! n

-- | Similarly, we can look up a list of indices in our DAG with repeated calls
-- to the above function. In fact, the two are mutually recursive!
(!=!) :: Monad f => DagF f xs -> HListF SFin idxs -> f (HList_ ((!!) xs) idxs)
(!=!) _    HNil     = pure HNil
(!=!) dag (i :& is) = (:&) <$> dag !! i <*> dag !=! is

-- | The above two functions mean that we can naïvely compute the top node of
-- our DAG simply by asking for the node at position zero.
evaluate :: Monad f => DagF f (x ':> xs) -> f x
evaluate = (!! SFZ)

-- | The above functions are fine, but we quite often /don't/ want to recompute
-- nodes in our DAG. For example, if the nodes' computations are relatively
-- expensive, we would like to calculate once and then reuse this result
-- whenever it is needed. To do this, we can create a "cache" for our DAG as a
-- list of /maybe/ the results at each node.
type Cache = HListF Maybe

-- | Also, in order to facilitate a cache, we need to know how to build empty
-- caches, and we'll thus need a pretty simple type class.
class                Empty (xs :: Vector Type n) where empty :: HListF Maybe xs
instance             Empty  'VNil                where empty = HNil
instance Empty xs => Empty (x ':> xs)            where empty = Nothing :& empty

-- | This function is identical to the '(!!)' function, except that it builds
-- up a cache of results already computed, and will lookup any node that needs
-- to be calculated within its cache preferably to re-computing.
(!!+) :: (Empty xs, Monad f) => DagF f xs -> SFin idx -> f (xs !!. idx)
(!!+) xs = fmap fst . getOneCached empty xs

-- | Internally, the '(!!+)' function calls here with an initially empty cache.
-- This function takes that cache, finds the desired node, and then calculates
-- (with the cache), its dependency tree. At that point, it calculates the
-- root. Of course, if the root already exists in the cache... well, we just
-- return that!
getOneCached
  :: Monad f
  => Cache xs
  -> DagF f xs
  -> SFin idx
  -> f (xs !!. idx, Cache xs)

getOneCached cache node (SFS n) -- Wrong node - keep recursing!
  = case (cache, node) of
      (c :& cs, Node _ _ xs) ->
        fmap (c :&) <$> getOneCached cs xs n

getOneCached cache@(Just x :& _) (Node _ _ _) SFZ -- Right node and cached!
  = pure (x, cache)

getOneCached (Nothing :& rest) (Node indices f tail) SFZ = do -- Cache miss
  (dependents, rest) <- getManyCached rest tail indices
  result             <- f dependents

  pure (result, Just result :& rest)

-- | Just as with '(!=!)', we can write '(!!+)' to work over a list of indices
-- with some mutual recursion (which does, indeed, share the cache!)
(!=!+) :: (Empty xs, Monad f) => DagF f xs -> HListF SFin idxs -> f (HList_ ((!!) xs) idxs)
(!=!+) xs = fmap fst . getManyCached empty xs

-- | As with 'getOneCached', we carry the cache with us, and compute a list of
-- nodes from the DAG. Mostly, this is just an internal function for
-- 'getOneCached'.
getManyCached
  :: Monad f
  => Cache xs
  -> DagF f xs
  -> HListF SFin indices
  -> f (HList_ ((!!) xs) indices, Cache xs)

getManyCached cached dag = \case
  HNil    -> pure (HNil, cached)
  i :& is -> do
    (head, cached) <- getOneCached  cached dag i
    (tail, cached) <- getManyCached cached dag is

    pure (head :& tail, cached)

-- | Evaluation is just accessing the top node again! The difference here is
-- that any nodes we see more than once will only be calculated once. Pretty
-- neat!
evaluate' :: (Empty xs, Monad f) => DagF f (x ':> xs) -> f x
evaluate' = (!!+ SFZ)