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
  81 -- Control --------------------------------------------------------------------
  82 -- | Return some result saying we've failed.
  83 noV :: VM a
  84 noV     = VM $ \_ _ _ -> return No
  85
  86
  87 -- | Like `traceNoV` but also emit some trace message to stderr.
  88 traceNoV :: String -> SDoc -> VM a
  89 traceNoV s d  = pprTrace s d noV
  90
  91
  92 -- | If `True` then carry on, otherwise fail.
  93 ensureV :: Bool -> VM ()
  94 ensureV False = noV
  95 ensureV True  = return ()
  96
  97
  98 -- | Like `ensureV` but if we fail then emit some trace message to stderr.
  99 traceEnsureV :: String -> SDoc -> Bool -> VM ()
 100 traceEnsureV s d False = traceNoV s d
 101 traceEnsureV _ _ True  = return ()
 102
 103
 104 -- | If `True` then return the first argument, otherwise fail.
 105 onlyIfV :: Bool -> VM a -> VM a
 106 onlyIfV b p = ensureV b >> p
 107
 108
 109 -- | Try some vectorisation computaton.
 110 --      If it succeeds then return `Just` the result,
 111 --      otherwise return `Nothing`.
 112 tryV :: VM a -> VM (Maybe a)
 113 tryV (VM p) = VM $ \bi genv lenv ->
 114   do
 115     r <- p bi genv lenv
 116     case r of
 117       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 118       No                -> return (Yes genv  lenv  Nothing)
 119
 120
 121 -- | If `Just` then return the value, otherwise fail.
 122 maybeV :: VM (Maybe a) -> VM a
 123 maybeV p = maybe noV return =<< p
 124
 125
 126 -- | Like `maybeV` but emit a message to stderr if we fail.
 127 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 128 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 129
 130
 131 -- | Try the first computation,
 132 --      if it succeeds then take the returned value,
 133 --      if it fails then run the second computation instead.
 134 orElseV :: VM a -> VM a -> VM a
 135 orElseV p q = maybe q return =<< tryV p
 136
 137
 138 -- | Fixpoint in the vectorisation monad.
 139 fixV :: (a -> VM a) -> VM a
 140 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 141   where
 142     -- NOTE: It is essential that we are lazy in r above so do not replace
 143     --       calls to this function by an explicit case.
 144     unYes (Yes _ _ x) = x
 145     unYes No          = panic "Vectorise.Monad.Base.fixV: no result"
 146