compiler/vectorise/Vectorise/Monad/Base.hs

   1
   2 -- | The Vectorisation monad.
   3 module Vectorise.Monad.Base (
   4         -- * The Vectorisation Monad
   5         VResult(..),
   6         VM(..),
   7
   8         -- * Lifting
   9         liftDs,
  10
  11         -- * Error Handling
  12         cantVectorise,
  13         maybeCantVectorise,
  14         maybeCantVectoriseM,
  15
  16         -- * Control
  17         noV,     traceNoV,
  18         ensureV, traceEnsureV,
  19         onlyIfV,
  20         tryV,
  21         maybeV,  traceMaybeV,
  22         orElseV,
  23         fixV,
  24 ) where
  25 import Vectorise.Builtins
  26 import Vectorise.Env
  27
  28 import DsMonad
  29 import Outputable
  30
  31
  32 -- The Vectorisation Monad ----------------------------------------------------
  33 -- | Vectorisation can either succeed with new envionment and a value,
  34 --   or return with failure.
  35 data VResult a
  36         = Yes GlobalEnv LocalEnv a | No
  37
  38 newtype VM a
  39         = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
  40
  41 instance Monad VM where
  42   return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
  43   VM p >>= f = VM $ \bi genv lenv -> do
  44                                       r <- p bi genv lenv
  45                                       case r of
  46                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
  47                                         No                -> return No
  48
  49
  50 -- Lifting --------------------------------------------------------------------
  51 -- | Lift a desugaring computation into the vectorisation monad.
  52 liftDs :: DsM a -> VM a
  53 liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
  54
  55
  56 -- Error Handling -------------------------------------------------------------
  57 -- | Throw a `pgmError` saying we can't vectorise something.
  58 cantVectorise :: String -> SDoc -> a
  59 cantVectorise s d = pgmError
  60                   . showSDocDump
  61                   $ vcat [text "*** Vectorisation error ***",
  62                           nest 4 $ sep [text s, nest 4 d]]
  63
  64
  65 -- | Like `fromJust`, but `pgmError` on Nothing.
  66 maybeCantVectorise :: String -> SDoc -> Maybe a -> a
  67 maybeCantVectorise s d Nothing  = cantVectorise s d
  68 maybeCantVectorise _ _ (Just x) = x
  69
  70
  71 -- | Like `maybeCantVectorise` but in a `Monad`.
  72 maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
  73 maybeCantVectoriseM s d p
  74   = do
  75       r <- p
  76       case r of
  77         Just x  -> return x
  78         Nothing -> cantVectorise s d
  79
  80 -- Control --------------------------------------------------------------------
  81 -- | Return some result saying we've failed.
  82 noV :: VM a
  83 noV     = VM $ \_ _ _ -> return No
  84
  85
  86 -- | Like `traceNoV` but also emit some trace message to stderr.
  87 traceNoV :: String -> SDoc -> VM a
  88 traceNoV s d  = pprTrace s d noV
  89
  90
  91 -- | If `True` then carry on, otherwise fail.
  92 ensureV :: Bool -> VM ()
  93 ensureV False = noV
  94 ensureV True  = return ()
  95
  96
  97 -- | Like `ensureV` but if we fail then emit some trace message to stderr.
  98 traceEnsureV :: String -> SDoc -> Bool -> VM ()
  99 traceEnsureV s d False = traceNoV s d
 100 traceEnsureV _ _ True  = return ()
 101
 102
 103 -- | If `True` then return the first argument, otherwise fail.
 104 onlyIfV :: Bool -> VM a -> VM a
 105 onlyIfV b p = ensureV b >> p
 106
 107
 108 -- | Try some vectorisation computaton.
 109 --      If it succeeds then return `Just` the result,
 110 --      otherwise return `Nothing`.
 111 tryV :: VM a -> VM (Maybe a)
 112 tryV (VM p) = VM $ \bi genv lenv ->
 113   do
 114     r <- p bi genv lenv
 115     case r of
 116       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 117       No                -> return (Yes genv  lenv  Nothing)
 118
 119
 120 -- | If `Just` then return the value, otherwise fail.
 121 maybeV :: VM (Maybe a) -> VM a
 122 maybeV p = maybe noV return =<< p
 123
 124
 125 -- | Like `maybeV` but emit a message to stderr if we fail.
 126 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 127 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 128
 129
 130 -- | Try the first computation,
 131 --      if it succeeds then take the returned value,
 132 --      if it fails then run the second computation instead.
 133 orElseV :: VM a -> VM a -> VM a
 134 orElseV p q = maybe q return =<< tryV p
 135
 136
 137 -- | Fixpoint in the vectorisation monad.
 138 fixV :: (a -> VM a) -> VM a
 139 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 140   where
 141     -- NOTE: It is essential that we are lazy in r above so do not replace
 142     --       calls to this function by an explicit case.
 143     unYes (Yes _ _ x) = x
 144     unYes No          = panic "Vectorise.Monad.Base.fixV: no result"
 145