Spark.hs

{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -feager-blackholing #-}

module Spark
  ( spark,
    sparkST,
    sparkList,
  )
where

import GHC.Exts
import GHC.ST

{-# INLINE spark #-}
spark :: a -> ST s a
spark t = ST $ \s0 -> case spark# (lazy t) s0 of
  (# s1, r #) -> (# s1, lazy r #)

-- Do not use noDuplicate# to prevent duplicate work! It's very
-- expensive and only suitable to be used sparingly where preventing
-- duplicate work is crucial for program logic! Sparks is meant to
-- handle thunks with pure computation that can tolerate duplicated
-- evaluation, so don't pay for what you don't use. On the other hand,
-- applying -feager-blackholing on this module and its use sites
-- should be good enough if you care, and it's relatively cheap,
-- though it's strongly advised to do your own benchmark to confirm
-- whether it's worthy for your use case.

{-# INLINE sparkST #-}
sparkST :: ST s a -> ST s a
sparkST m = spark $ case runRW# $ unsafeCoerce# $ lazy m of
  (# _, r #) -> lazy r

sparkList :: [a] -> ST s [a]
sparkList l = sparkST $ case l of
  [] -> pure []
  x : xs -> do
    x' <- spark x
    xs' <- sparkList xs
    pure $ x' : xs'