Flatten.hs

   1 module Flatten where
   2 import CoreSyn
   3 import Control.Monad
   4 import qualified Type
   5 import qualified Name
   6 import qualified TyCon
   7 import qualified Maybe
   8 import Data.Traversable
   9 import qualified DataCon
  10 import qualified CoreUtils
  11 import Control.Applicative
  12 import Outputable ( showSDoc, ppr )
  13 import qualified Data.Foldable as Foldable
  14 import qualified Control.Monad.State as State
  15
  16 -- | A datatype that maps each of the single values in a haskell structure to
  17 -- a mapto. The map has the same structure as the haskell type mapped, ie
  18 -- nested tuples etc.
  19 data HsValueMap mapto =
  20   Tuple [HsValueMap mapto]
  21   | Single mapto
  22   deriving (Show, Eq)
  23
  24 instance Functor HsValueMap where
  25   fmap f (Single s) = Single (f s)
  26   fmap f (Tuple maps) = Tuple (map (fmap f) maps)
  27
  28 instance Foldable.Foldable HsValueMap where
  29   foldMap f (Single s) = f s
  30   -- The first foldMap folds a list of HsValueMaps, the second foldMap folds
  31   -- each of the HsValueMaps in that list
  32   foldMap f (Tuple maps) = Foldable.foldMap (Foldable.foldMap f) maps
  33
  34 instance Traversable HsValueMap where
  35   traverse f (Single s) = Single <$> f s
  36   traverse f (Tuple maps) = Tuple <$> (traverse (traverse f) maps)
  37
  38 data PassState s x = PassState (s -> (s, x))
  39
  40 instance Functor (PassState s) where
  41   fmap f (PassState a) = PassState (\s -> let (s', a') = a s in (s', f a'))
  42
  43 instance Applicative (PassState s) where
  44   pure x = PassState (\s -> (s, x))
  45   PassState f <*> PassState x = PassState (\s -> let (s', f') = f s; (s'', x') = x s' in (s'', f' x'))
  46
  47 -- | Creates a HsValueMap with the same structure as the given type, using the
  48 --   given function for mapping the single types.
  49 mkHsValueMap ::
  50   Type.Type                         -- ^ The type to map to a HsValueMap
  51   -> HsValueMap Type.Type           -- ^ The resulting map and state
  52
  53 mkHsValueMap ty =
  54   case Type.splitTyConApp_maybe ty of
  55     Just (tycon, args) ->
  56       if (TyCon.isTupleTyCon tycon)
  57         then
  58           Tuple (map mkHsValueMap args)
  59         else
  60           Single ty
  61     Nothing -> Single ty
  62
  63 -- Extract the arguments from a data constructor application (that is, the
  64 -- normal args, leaving out the type args).
  65 dataConAppArgs :: DataCon.DataCon -> [CoreExpr] -> [CoreExpr]
  66 dataConAppArgs dc args =
  67     drop tycount args
  68   where
  69     tycount = length $ DataCon.dataConAllTyVars dc
  70
  71
  72
  73 data FlatFunction = FlatFunction {
  74   args   :: [SignalDefMap],
  75   res    :: SignalUseMap,
  76   --sigs   :: [SignalDef],
  77   apps   :: [FApp],
  78   conds  :: [CondDef]
  79 } deriving (Show, Eq)
  80
  81 type SignalUseMap = HsValueMap SignalUse
  82 type SignalDefMap = HsValueMap SignalDef
  83
  84 useMapToDefMap :: SignalUseMap -> SignalDefMap
  85 useMapToDefMap = fmap (\(SignalUse u) -> SignalDef u)
  86
  87 defMapToUseMap :: SignalDefMap -> SignalUseMap
  88 defMapToUseMap = fmap (\(SignalDef u) -> SignalUse u)
  89
  90
  91 type SignalId = Int
  92 data SignalUse = SignalUse {
  93   sigUseId :: SignalId
  94 } deriving (Show, Eq)
  95
  96 data SignalDef = SignalDef {
  97   sigDefId :: SignalId
  98 } deriving (Show, Eq)
  99
 100 data FApp = FApp {
 101   appFunc :: HsFunction,
 102   appArgs :: [SignalUseMap],
 103   appRes  :: SignalDefMap
 104 } deriving (Show, Eq)
 105
 106 data CondDef = CondDef {
 107   cond    :: SignalUse,
 108   high    :: SignalUse,
 109   low     :: SignalUse,
 110   condRes :: SignalDef
 111 } deriving (Show, Eq)
 112
 113 -- | How is a given (single) value in a function's type (ie, argument or
 114 -- return value) used?
 115 data HsValueUse =
 116   Port           -- ^ Use it as a port (input or output)
 117   | State Int    -- ^ Use it as state (input or output). The int is used to
 118                  --   match input state to output state.
 119   | HighOrder {  -- ^ Use it as a high order function input
 120     hoName :: String,  -- ^ Which function is passed in?
 121     hoArgs :: [HsUseMap]   -- ^ Which arguments are already applied? This
 122                          -- ^ map should only contain Port and other
 123                          --   HighOrder values.
 124   }
 125   deriving (Show, Eq)
 126
 127 type HsUseMap = HsValueMap HsValueUse
 128
 129 -- | Builds a HsUseMap with the same structure has the given HsValueMap in
 130 --   which all the Single elements are marked as State, with increasing state
 131 --   numbers.
 132 useAsState :: HsValueMap a -> HsUseMap
 133 useAsState map =
 134   map'
 135   where
 136     -- Traverse the existing map, resulting in a function that maps an initial
 137     -- state number to the final state number and the new map
 138     PassState f = traverse asState map
 139     -- Run this function to get the new map
 140     (_, map')   = f 0
 141     -- This function maps each element to a State with a unique number, by
 142     -- incrementing the state count.
 143     asState x   = PassState (\s -> (s+1, State s))
 144
 145 -- | Builds a HsUseMap with the same structure has the given HsValueMap in
 146 --   which all the Single elements are marked as Port.
 147 useAsPort :: HsValueMap a -> HsUseMap
 148 useAsPort map = fmap (\x -> Port) map
 149
 150 data HsFunction = HsFunction {
 151   hsFuncName :: String,
 152   hsFuncArgs :: [HsUseMap],
 153   hsFuncRes  :: HsUseMap
 154 } deriving (Show, Eq)
 155
 156 type BindMap = [(
 157   CoreBndr,            -- ^ The bind name
 158   Either               -- ^ The bind value which is either
 159     SignalUseMap       -- ^ a signal
 160     (
 161       HsValueUse,      -- ^ or a HighOrder function
 162       [SignalUse]      -- ^ With these signals already applied to it
 163     )
 164   )]
 165
 166 type FlattenState = State.State ([FApp], [CondDef], SignalId)
 167
 168 -- | Add an application to the current FlattenState
 169 addApp :: FApp -> FlattenState ()
 170 addApp a = do
 171   (apps, conds, n) <- State.get
 172   State.put (a:apps, conds, n)
 173
 174 -- | Add a conditional definition to the current FlattenState
 175 addCondDef :: CondDef -> FlattenState ()
 176 addCondDef c = do
 177   (apps, conds, n) <- State.get
 178   State.put (apps, c:conds, n)
 179
 180 -- | Generates a new signal id, which is unique within the current flattening.
 181 genSignalId :: FlattenState SignalId
 182 genSignalId = do
 183   (apps, conds, n) <- State.get
 184   State.put (apps, conds, n+1)
 185   return n
 186
 187 genSignalUses ::
 188   Type.Type
 189   -> FlattenState SignalUseMap
 190
 191 genSignalUses ty = do
 192   typeMapToUseMap tymap
 193   where
 194     -- First generate a map with the right structure containing the types
 195     tymap = mkHsValueMap ty
 196
 197 typeMapToUseMap ::
 198   HsValueMap Type.Type
 199   -> FlattenState SignalUseMap
 200
 201 typeMapToUseMap (Single ty) = do
 202   id <- genSignalId
 203   return $ Single (SignalUse id)
 204
 205 typeMapToUseMap (Tuple tymaps) = do
 206   usemaps <- State.mapM typeMapToUseMap tymaps
 207   return $ Tuple usemaps
 208
 209 -- | Flatten a haskell function
 210 flattenFunction ::
 211   HsFunction                      -- ^ The function to flatten
 212   -> CoreBind                     -- ^ The function value
 213   -> FlatFunction                 -- ^ The resulting flat function
 214
 215 flattenFunction _ (Rec _) = error "Recursive binders not supported"
 216 flattenFunction hsfunc bind@(NonRec var expr) =
 217   FlatFunction args res apps conds
 218   where
 219     init_state        = ([], [], 0)
 220     (fres, end_state) = State.runState (flattenExpr [] expr) init_state
 221     (args, res)       = fres
 222     (apps, conds, _)  = end_state
 223
 224 flattenExpr ::
 225   BindMap
 226   -> CoreExpr
 227   -> FlattenState ([SignalDefMap], SignalUseMap)
 228
 229 flattenExpr binds lam@(Lam b expr) = do
 230   -- Find the type of the binder
 231   let (arg_ty, _) = Type.splitFunTy (CoreUtils.exprType lam)
 232   -- Create signal names for the binder
 233   defs <- genSignalUses arg_ty
 234   let binds' = (b, Left defs):binds
 235   (args, res) <- flattenExpr binds' expr
 236   return ((useMapToDefMap defs) : args, res)
 237
 238 flattenExpr binds (Var id) =
 239   case bind of
 240     Left sig_use -> return ([], sig_use)
 241     Right _ -> error "Higher order functions not supported."
 242   where
 243     bind = Maybe.fromMaybe
 244       (error $ "Argument " ++ Name.getOccString id ++ "is unknown")
 245       (lookup id binds)
 246
 247 flattenExpr binds app@(App _ _) = do
 248   -- Is this a data constructor application?
 249   case CoreUtils.exprIsConApp_maybe app of
 250     -- Is this a tuple construction?
 251     Just (dc, args) -> if DataCon.isTupleCon dc
 252       then
 253         flattenBuildTupleExpr binds (dataConAppArgs dc args)
 254       else
 255         error $ "Data constructors other than tuples not supported: " ++ (showSDoc $ ppr app)
 256     otherwise ->
 257       -- Normal function application
 258       let ((Var f), args) = collectArgs app in
 259       flattenApplicationExpr binds (CoreUtils.exprType app) f args
 260   where
 261     flattenBuildTupleExpr = error $ "Tuple construction not supported: " ++ (showSDoc $ ppr app)
 262     flattenApplicationExpr binds ty f args = error $ "Function application not supported: " ++ (showSDoc $ ppr app)
 263
 264 flattenExpr _ _ = do
 265   return ([], Tuple [])
 266
 267
 268 -- vim: set ts=8 sw=2 sts=2 expandtab: