module Hoopl (
    module Compiler.Hoopl,
    module Hoopl.Dataflow,
    deepBwdRw3, deepBwdRw,
    thenFwdRw
  ) where

import Compiler.Hoopl hiding
  ( Unique,
    FwdTransfer(..), FwdRewrite(..), FwdPass(..),
    BwdTransfer(..), BwdRewrite(..), BwdPass(..),
    noFwdRewrite, noBwdRewrite,
--    analyzeAndRewriteFwd, analyzeAndRewriteBwd,
    mkFactBase, Fact,
    mkBRewrite, mkBRewrite3, mkBTransfer, mkBTransfer3,
    mkFRewrite, mkFRewrite3, mkFTransfer, mkFTransfer3,
    deepFwdRw, deepFwdRw3, thenFwdRw, iterFwdRw,
    deepBwdRw, deepBwdRw3, thenBwdRw, iterBwdRw
  )

import Hoopl.Dataflow
import OptimizationFuel
import Control.Monad

deepFwdRw3 :: (n C O -> f -> FuelUniqSM (Maybe (Graph n C O)))
           -> (n O O -> f -> FuelUniqSM (Maybe (Graph n O O)))
           -> (n O C -> f -> FuelUniqSM (Maybe (Graph n O C)))
           -> (FwdRewrite FuelUniqSM n f)
deepFwdRw :: (forall e x . n e x -> f -> FuelUniqSM (Maybe (Graph n e x))) -> FwdRewrite FuelUniqSM n f
deepFwdRw3 f m l = iterFwdRw $ mkFRewrite3 f m l
deepFwdRw f = deepFwdRw3 f f f

-- N.B. rw3, rw3', and rw3a are triples of functions.
-- But rw and rw' are single functions.
thenFwdRw :: forall n f.
             FwdRewrite FuelUniqSM n f
          -> FwdRewrite FuelUniqSM n f 
          -> FwdRewrite FuelUniqSM n f
thenFwdRw rw3 rw3' = wrapFR2 thenrw rw3 rw3'
 where
  thenrw :: forall e x t t1.
               (t -> t1 -> FuelUniqSM (Maybe (Graph n e x, FwdRewrite FuelUniqSM n f)))
            -> (t -> t1 -> FuelUniqSM (Maybe (Graph n e x, FwdRewrite FuelUniqSM n f)))
            -> t
            -> t1
            -> FuelUniqSM (Maybe (Graph n e x, FwdRewrite FuelUniqSM n f))
  thenrw rw rw' n f = rw n f >>= fwdRes
     where fwdRes Nothing   = rw' n f
           fwdRes (Just gr) = return $ Just $ fadd_rw rw3' gr

iterFwdRw :: forall m n f.
             FwdRewrite FuelUniqSM n f
          -> FwdRewrite FuelUniqSM n f
iterFwdRw rw3 = wrapFR iter rw3
 where iter :: forall a e x t.
               (t -> a -> FuelUniqSM (Maybe (Graph n e x, FwdRewrite FuelUniqSM n f)))
               -> t
               -> a
               -> FuelUniqSM (Maybe (Graph n e x, FwdRewrite FuelUniqSM n f))
       iter rw n = (liftM $ liftM $ fadd_rw (iterFwdRw rw3)) . rw n

-- | Function inspired by 'rew' in the paper
_frewrite_cps :: ((Graph n e x, FwdRewrite FuelUniqSM n f) -> FuelUniqSM a)
             -> FuelUniqSM a
             -> (forall e x . n e x -> f -> FuelUniqSM (Maybe (Graph n e x, FwdRewrite FuelUniqSM n f)))
             -> n e x
             -> f
             -> FuelUniqSM a
_frewrite_cps j n rw node f =
    do mg <- rw node f
       case mg of Nothing -> n
                  Just gr -> j gr



-- | Function inspired by 'add' in the paper
fadd_rw :: FwdRewrite FuelUniqSM n f
        -> (Graph n e x, FwdRewrite FuelUniqSM n f)
        -> (Graph n e x, FwdRewrite FuelUniqSM n f)
fadd_rw rw2 (g, rw1) = (g, rw1 `thenFwdRw` rw2)



deepBwdRw3 ::
              (n C O -> f          -> FuelUniqSM (Maybe (Graph n C O)))
           -> (n O O -> f          -> FuelUniqSM (Maybe (Graph n O O)))
           -> (n O C -> FactBase f -> FuelUniqSM (Maybe (Graph n O C)))
           -> (BwdRewrite FuelUniqSM n f)
deepBwdRw  :: (forall e x . n e x -> Fact x f -> FuelUniqSM (Maybe (Graph n e x)))
           -> BwdRewrite FuelUniqSM n f
deepBwdRw3 f m l = iterBwdRw $ mkBRewrite3 f m l
deepBwdRw  f = deepBwdRw3 f f f


thenBwdRw :: forall n f. BwdRewrite FuelUniqSM n f -> BwdRewrite FuelUniqSM n f -> BwdRewrite FuelUniqSM n f
thenBwdRw rw1 rw2 = wrapBR2 f rw1 rw2
  where f :: forall t t1 t2 e x.
             t
             -> (t1 -> t2 -> FuelUniqSM (Maybe (Graph n e x, BwdRewrite FuelUniqSM n f)))
             -> (t1 -> t2 -> FuelUniqSM (Maybe (Graph n e x, BwdRewrite FuelUniqSM n f)))
             -> t1
             -> t2
             -> FuelUniqSM (Maybe (Graph n e x, BwdRewrite FuelUniqSM n f))
        f _ rw1 rw2' n f = do
          res1 <- rw1 n f
          case res1 of
            Nothing -> rw2' n f
            Just gr -> return $ Just $ badd_rw rw2 gr

iterBwdRw :: forall n f. BwdRewrite FuelUniqSM n f -> BwdRewrite FuelUniqSM n f
iterBwdRw rw = wrapBR f rw
  where f :: forall t e x t1 t2.
             t
             -> (t1 -> t2 -> FuelUniqSM (Maybe (Graph n e x, BwdRewrite FuelUniqSM n f)))
             -> t1
             -> t2
             -> FuelUniqSM (Maybe (Graph n e x, BwdRewrite FuelUniqSM n f))
        f _ rw' n f = liftM (liftM (badd_rw (iterBwdRw rw))) (rw' n f)

-- | Function inspired by 'add' in the paper
badd_rw :: BwdRewrite FuelUniqSM n f
        -> (Graph n e x, BwdRewrite FuelUniqSM n f)
        -> (Graph n e x, BwdRewrite FuelUniqSM n f)
badd_rw rw2 (g, rw1) = (g, rw1 `thenBwdRw` rw2)