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         -- * Debugging
  17         traceVt, dumpOptVt, dumpVt,
  18
  19         -- * Control
  20         noV,     traceNoV,
  21         ensureV, traceEnsureV,
  22         onlyIfV,
  23         tryV,
  24         maybeV,  traceMaybeV,
  25         orElseV,
  26         fixV,
  27 ) where
  28
  29 import Vectorise.Builtins
  30 import Vectorise.Env
  31
  32 import DsMonad
  33 import TcRnMonad
  34 import ErrUtils
  35 import Outputable
  36 import DynFlags
  37 import StaticFlags
  38
  39 import Control.Monad
  40 import System.IO (stderr)
  41
  42
  43 -- The Vectorisation Monad ----------------------------------------------------
  44
  45 -- | Vectorisation can either succeed with new envionment and a value,
  46 --   or return with failure.
  47 data VResult a
  48         = Yes GlobalEnv LocalEnv a | No
  49
  50 newtype VM a
  51         = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
  52
  53 instance Monad VM where
  54   return x   = VM $ \_  genv lenv -> return (Yes genv lenv x)
  55   VM p >>= f = VM $ \bi genv lenv -> do
  56                                       r <- p bi genv lenv
  57                                       case r of
  58                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
  59                                         No                -> return No
  60
  61 instance Functor VM where
  62   fmap = liftM
  63
  64 instance MonadIO VM where
  65   liftIO = liftDs . liftIO
  66
  67
  68 -- Lifting --------------------------------------------------------------------
  69 -- | Lift a desugaring computation into the vectorisation monad.
  70 liftDs :: DsM a -> VM a
  71 liftDs p = VM $ \_ genv lenv -> do { x <- p; return (Yes genv lenv x) }
  72
  73
  74 -- Error Handling -------------------------------------------------------------
  75 -- | Throw a `pgmError` saying we can't vectorise something.
  76 cantVectorise :: String -> SDoc -> a
  77 cantVectorise s d = pgmError
  78                   . showSDocDump
  79                   $ vcat [text "*** Vectorisation error ***",
  80                           nest 4 $ sep [text s, nest 4 d]]
  81
  82
  83 -- | Like `fromJust`, but `pgmError` on Nothing.
  84 maybeCantVectorise :: String -> SDoc -> Maybe a -> a
  85 maybeCantVectorise s d Nothing  = cantVectorise s d
  86 maybeCantVectorise _ _ (Just x) = x
  87
  88
  89 -- | Like `maybeCantVectorise` but in a `Monad`.
  90 maybeCantVectoriseM :: Monad m => String -> SDoc -> m (Maybe a) -> m a
  91 maybeCantVectoriseM s d p
  92   = do
  93       r <- p
  94       case r of
  95         Just x  -> return x
  96         Nothing -> cantVectorise s d
  97
  98
  99 -- Debugging ------------------------------------------------------------------
 100
 101 -- |Output a trace message if -ddump-vt-trace is active.
 102 --
 103 traceVt :: String -> SDoc -> VM ()
 104 traceVt herald doc
 105   | 1 <= opt_TraceLevel = liftDs $
 106                             traceOptIf Opt_D_dump_vt_trace $
 107                               hang (text herald) 2 doc
 108   | otherwise           = return ()
 109
 110 -- |Dump the given program conditionally.
 111 --
 112 dumpOptVt :: DynFlag -> String -> SDoc -> VM ()
 113 dumpOptVt flag header doc
 114   = do { b <- liftDs $ doptM flag
 115        ; if b
 116          then dumpVt header doc
 117          else return ()
 118        }
 119
 120 -- |Dump the given program unconditionally.
 121 --
 122 dumpVt :: String -> SDoc -> VM ()
 123 dumpVt header doc
 124   = do { unqual <- liftDs mkPrintUnqualifiedDs
 125        ; liftIO $ printForUser stderr unqual (mkDumpDoc header doc)
 126        }
 127
 128 -- Control --------------------------------------------------------------------
 129 -- | Return some result saying we've failed.
 130 noV :: VM a
 131 noV     = VM $ \_ _ _ -> return No
 132
 133
 134 -- | Like `traceNoV` but also emit some trace message to stderr.
 135 traceNoV :: String -> SDoc -> VM a
 136 traceNoV s d  = pprTrace s d noV
 137
 138
 139 -- | If `True` then carry on, otherwise fail.
 140 ensureV :: Bool -> VM ()
 141 ensureV False = noV
 142 ensureV True  = return ()
 143
 144
 145 -- | Like `ensureV` but if we fail then emit some trace message to stderr.
 146 traceEnsureV :: String -> SDoc -> Bool -> VM ()
 147 traceEnsureV s d False = traceNoV s d
 148 traceEnsureV _ _ True  = return ()
 149
 150
 151 -- | If `True` then return the first argument, otherwise fail.
 152 onlyIfV :: Bool -> VM a -> VM a
 153 onlyIfV b p = ensureV b >> p
 154
 155
 156 -- | Try some vectorisation computaton.
 157 --      If it succeeds then return `Just` the result,
 158 --      otherwise return `Nothing`.
 159 tryV :: VM a -> VM (Maybe a)
 160 tryV (VM p) = VM $ \bi genv lenv ->
 161   do
 162     r <- p bi genv lenv
 163     case r of
 164       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 165       No                -> return (Yes genv  lenv  Nothing)
 166
 167
 168 -- | If `Just` then return the value, otherwise fail.
 169 maybeV :: VM (Maybe a) -> VM a
 170 maybeV p = maybe noV return =<< p
 171
 172
 173 -- | Like `maybeV` but emit a message to stderr if we fail.
 174 traceMaybeV :: String -> SDoc -> VM (Maybe a) -> VM a
 175 traceMaybeV s d p = maybe (traceNoV s d) return =<< p
 176
 177
 178 -- | Try the first computation,
 179 --      if it succeeds then take the returned value,
 180 --      if it fails then run the second computation instead.
 181 orElseV :: VM a -> VM a -> VM a
 182 orElseV p q = maybe q return =<< tryV p
 183
 184
 185 -- | Fixpoint in the vectorisation monad.
 186 fixV :: (a -> VM a) -> VM a
 187 fixV f = VM (\bi genv lenv -> fixDs $ \r -> runVM (f (unYes r)) bi genv lenv )
 188   where
 189     -- NOTE: It is essential that we are lazy in r above so do not replace
 190     --       calls to this function by an explicit case.
 191     unYes (Yes _ _ x) = x
 192     unYes No          = panic "Vectorise.Monad.Base.fixV: no result"
 193