mldsa-0.1.0.0: src/Fusion.hs

-- |
-- Module      : Fusion
-- License     : BSD-3-Clause
-- Copyright   : (c) 2026 Olivier Chéron
--
-- Infrastructure to decrease intermediate allocations and prefer in-place
-- mutation when possible
--
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilyDependencies #-}
module Fusion
    ( Fusion(..), MapF(..)
    , Context, runContext, newContext, thawContext, mapContext, modifyContext
    , foldContext, seqContext
    ) where

import Control.Monad ( forM_, (>=>) )
import Control.Monad.ST

-- class of values that can be mutated in the ST monad
class Fusion a where
    type Mut a s = mut | mut -> a
    newF :: ST s (Mut a s)
    thawF :: a -> ST s (Mut a s)
    unsafeFreezeF :: Mut a s -> ST s a

-- a transformation step in the fusion pipeline, with two implementations
-- provided: one that operates on an existing mutation context, and one that
-- initiates a new context from the input
data MapF a b = MapF
    { mapUpdate :: forall s. Mut a s -> ST s (Mut b s)
    , mapInit :: forall s. a -> ST s (Mut b s)
    }

-- MapF is almost a category except for the 'Fusion' constraint on objects
--
-- idMapF :: Fusion a => MapF a a
-- idMapF = MapF { mapUpdate = pure, mapInit = thawF }

composeMapF :: MapF b c -> MapF a b -> MapF a c
composeMapF m2 m1 = MapF
    { mapUpdate = mapUpdate m1 >=> mapUpdate m2
    , mapInit = mapInit m1 >=> mapUpdate m2
    }

-- fusion context
newtype Context a = Context (forall s. ST s (Mut a s))

newContext :: Fusion a => Context a
newContext = Context newF

thawContext :: Fusion a => a -> Context a
thawContext a = Context $ thawF a
{-# INLINE [0] thawContext #-}

modifyContext :: (forall s. Mut a s -> ST s ()) -> Context a -> Context a
modifyContext f = bindContext $ \ma -> f ma >> return ma

mapContext :: MapF a b -> Context a -> Context b
mapContext m = bindContext (mapUpdate m)
{-# INLINE [0] mapContext #-}

initContext :: MapF a b -> a -> Context b
initContext m a = Context $ mapInit m a

bindContext :: (forall s. Mut a s -> ST s (Mut b s)) -> Context a -> Context b
bindContext f (Context ctx) = Context $ ctx >>= f

foldContext :: Foldable t => (forall s. b -> Mut a s -> ST s ()) -> Context a -> t b -> Context a
foldContext f c bs = modifyContext (\ma -> forM_ bs $ \b -> f b ma) c

runContext :: Fusion a => Context a -> a
runContext (Context ctx) = runST (ctx >>= unsafeFreezeF)
{-# INLINE [0] runContext #-}

seqContext :: a -> Context b -> Context b
seqContext = seq
{-# INLINE [0] seqContext #-}


-- Fusion rules
--
-- "thawContext/runContext" is the canonical optimization that eliminates an
-- allocation + value copy.  Instead, it sequences two transformations on the
-- same mutation context.
--
-- "mapContext/seqContext" moves strictness annotations upstream so that they
-- do not prevent other rules from firing.
--
-- "mapContext/mapContext" is not strictly needed: the function is ultimately
-- inlined to the same code.  But we keep it so that simplifications fire early
-- and do not wait for the final phase.
--
-- "mapContext/thawContext" is the rule that invokes mapInit instead of copying
-- the input and calling mapUpdate.

{-# RULES
"thawContext/runContext" [~0] forall c. thawContext (runContext c) = c
"mapContext/seqContext" [~0] forall a m c. mapContext m (seqContext a c) = seqContext a (mapContext m c)
"mapContext/mapContext" [~0] forall m1 m2 c. mapContext m2 (mapContext m1 c) = mapContext (composeMapF m2 m1) c
"mapContext/thawContext" [1] forall m a. mapContext m (thawContext a) = initContext m a
  #-}