CodeGen.hs 13.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
-- |
-- Module      :  Compiler.Mu.CodeGen
-- Copyright   :  nathyong 2016
-- License     :  BSD3
--
-- Maintainer  :  nathyong@gmail.com
-- Stability   :  experimental
-- Portability :  unknown
--
-- Utility functions for Haskell code generation via Mu.
--

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
14 15 16
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE TemplateHaskell #-}
17

18 19 20 21 22
module Compiler.Mu.CodeGen
  ( -- * Mu Monad
    Mu
  , runMu
  , definitions
23
  , topLevels
24 25
  , currentBlock
  , MuResult
26 27
  , MuMergedResult
  , mergeResults
28 29 30
    -- * Block building monad
  , Blocks
  , runBlocks
31 32 33
  , basicBlock
  , assign
  , upValue
34
    -- * Code generation
35
  , UnboxedData (..)
36
  , bundleMu
37
  , populate
38
    -- * Mu interface
39 40
  , typedef
  , constant
41 42
  , funcsig
  , funcdef
43
  , primfuncdef
44 45
  , closure
  , string
46 47
    -- * Naming interface
  , dataConEntry
48 49 50 51 52 53 54 55
  , closureNameOf
  , entryNameOf
    -- * Type names
  , loadTypes
  , i64
  , i8
  , muString
    -- * STG Closures in Mu
56
  , closureTR
57 58 59 60 61
  , muClosure
  , muClosureRef
  , muClosureFunction
  , muClosureFunctionSig
  , haskellData
62
  , getClosureField
63
  , getClosureData
64
  , ClosureField (..)
65 66
) where

67
import Control.Monad.Trans (MonadTrans, lift)
68
import Control.Applicative ((<|>))
69 70
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.State (State, StateT, MonadState, evalState, runStateT, gets)
71
import Control.Monad.Writer (Writer, WriterT, MonadWriter, runWriterT, tell)
72
import Data.Binary (Binary)
73
import Data.Foldable (toList, foldl')
74
import Data.List (genericLength)
75
import Data.Map.Strict (Map)
76 77 78 79
import Data.Monoid ((<>))
import Data.Sequence (Seq, (|>))
import Data.String (fromString)
import GHC.Generics (Generic)
80
import qualified Data.ByteString as B
81
import qualified Data.Map.Strict as M
82 83
import qualified Foreign as F
import qualified Foreign.C as C
84

85
import           Lens.Micro.Platform ((%=), makeLenses)
86

87
import           FastString (fsLit, zEncodeFS, zString)
88
import qualified GHC
89
import           Name (NamedThing (..), nameStableString, isSystemName, isInternalName)
90
import           Unique (Uniquable (..))
91

92 93 94
import           Mu.API
import           Mu.AST
import           Mu.Interface
95

96
import Debug.Trace
97

98
-------------------------------------------------- * Data types
99

100 101 102 103 104 105 106
-- | Top-level constructs which need to be allocated and populated in the Mu
-- separately.
data TopLevel = Closure GlobalCellName [UnboxedData] FunctionName
              | ByteArray GlobalCellName B.ByteString
    deriving (Generic)

instance Binary TopLevel
107

108 109 110 111 112
instance HasName TopLevel where
    toName tl = toName $ case tl of
        Closure n _ _ -> n
        ByteArray n _ -> n

113 114 115
data UnboxedData = UnboxedInt Integer
                 | UnboxedDouble Rational
                 | TaggedPointer GlobalCellName Int
116 117
    deriving (Generic)

118
instance Binary UnboxedData
119

120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
type MuResult = (Seq Definition, Seq TopLevel, Maybe FunctionName)
type MuMergedResult = ( Map Name Definition
                      , Map Name TopLevel
                      , Maybe FunctionName)

mergeResults :: [MuResult] -> MuMergedResult
mergeResults = foldl' mergeResults' (M.empty, M.empty, Nothing)
  where
    mergeResults' :: MuMergedResult -> MuResult -> MuMergedResult
    mergeResults' (defmap, topmap, fn) (defs, tops, fn') =
        ( foldl' insert defmap defs
        , foldl' insert topmap tops
        , fn <|> fn')
      where
        insert m n = M.insert (toName n) n m
135 136 137 138 139 140

-------------------------------------------------- * Mu Monad

newtype Mu a = Mu (State MuState a)
    deriving (Functor, Applicative, Monad, MonadState MuState)

141 142
-- | Run a Mu computation from the empty state, storing the result.
runMu :: Mu a -> a
143
runMu (Mu code) = evalState code $ MuState 0 mempty mempty ()
144 145

data MuState = MuState
146
    { _uniqueID :: !Int
147 148
    , _definitions :: Seq Definition
    , _topLevels :: Seq TopLevel
149 150 151 152 153 154 155 156
    , _currentBlock :: ()
    }

makeLenses ''MuState

-------------------------------------------------- * Block building monad

data BlockState = BlockState { _blocks :: Seq BasicBlock }
157

158 159 160 161 162
makeLenses ''BlockState

newtype Blocks a = Blocks (StateT BlockState Mu a)
    deriving (Functor, Applicative, Monad, MonadState BlockState)

163 164 165
type BlocksWriter = WriterT [Assigned Expression] Blocks


166 167 168
runBlocks :: Blocks a -> Mu (a, Seq BasicBlock)
runBlocks (Blocks code) = do
    (a, bstate) <- runStateT code $ BlockState mempty
169
    -- TODO: allocate upvalues
170
    return (a, _blocks bstate)
171

172

173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207
-- | Lift a 'Mu' computation into the Blocks monad.
liftBlocks :: Mu a -> Blocks a
liftBlocks m = Blocks (lift m)


-- | Emit a basic block without an exceptional parameter.
basicBlock
    :: BasicBlockName
    -> [(VarName, TypedefName)]
    -> BlocksWriter Expression
    -> Blocks BasicBlock
basicBlock n params exprs = do
    (terminator, body) <- runWriterT exprs
    return $ BasicBlock n params Nothing body terminator
  where


-- | Assign the result of a Mu 'Expression' to a 'VarName'.
assign :: Expression -> BlocksWriter VarName
assign expr = do
    uid <- lift $ liftBlocks nextMuUnique
    let n = fromString ("var" ++ show uid)
    tell $ [[n] := expr]
    return n


-- | Declare a (constant) value to be defined at the top level.
upValue :: ConstConstructor -> TypedefName -> BlocksWriter ConstantName
upValue ctor ty = lift . liftBlocks $ do
    uid <- nextMuUnique
    let n = fromString ("@upvar" ++ show uid)
    constant n ty ctor
    return n


208 209
-------------------------------------------------- * Code generation

210
-- | Extract the definitions and other useful things from the Mu monad.
211 212 213 214
bundleMu :: Mu () -> MuResult
bundleMu codegen = runMu $ do
    codegen
    defns <- gets _definitions
215
    tops <- gets _topLevels
216
    return (defns, tops, Nothing)
217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232


-- | Populate a Mu Closure, or other top-level data type, at build-time.
-- Assumes that all the objects in a bundle will have been imported already.
populate :: MonadIO m => F.Ptr MuCtx -> TopLevel -> m ()
populate ctx top = case top of
    Closure n payload entry -> do
        -- allocate the closure
        closureLength <- handleFromSint64 ctx (genericLength payload) 64
        closureType <- getID ctx muClosure
        closureRef <- newHybrid ctx closureType closureLength

        -- populate the closure
        closureIref <- getIref ctx closureRef
        one <- handleFromSint64 ctx 1 52

233
        closureEntryRef <- getFieldIref ctx closureIref (fieldIx EntryField)
234 235 236
        entryRef <- getID ctx entry >>= handleFromFunc ctx
        store ctx muOrdNotAtomic closureEntryRef entryRef

237 238 239 240 241 242
        closureLengthRef <- getFieldIref ctx closureIref (fieldIx NArgsField)
        store ctx muOrdNotAtomic closureLengthRef closureLength

        closureTagRef <- getFieldIref ctx closureIref (fieldIx TagField)
        store ctx muOrdNotAtomic closureTagRef one

243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285
        closureDataRef <- getVarPartIref ctx closureIref
        let storePayloads datums ref = case datums of
                [] -> return ()
                d:ds -> do
                    case d of
                        UnboxedInt i -> do
                            val <- handleFromSint64 ctx (fromIntegral i) 52
                            trVal <- tr64FromInt ctx val
                            store ctx muOrdNotAtomic ref trVal
                        UnboxedDouble r -> do
                            val <- handleFromDouble ctx (fromRational r)
                            trVal <- tr64FromFp ctx val
                            store ctx muOrdNotAtomic ref trVal
                        TaggedPointer cell tag -> do
                            tag' <- handleFromSint64 ctx (fromIntegral tag) 6
                            cellID <- getID ctx cell
                            cellHandle <- handleFromGlobal ctx cellID
                            cellRef <- ctxLoad ctx muOrdNotAtomic cellHandle
                            trCellRef <- tr64FromRef ctx cellRef tag'
                            store ctx muOrdNotAtomic ref trCellRef
                    shiftIref ctx ref one >>= storePayloads ds

        storePayloads payload closureDataRef

        -- attach it to the global cell
        cell <- handleFromGlobal ctx =<< getID ctx n
        store ctx muOrdNotAtomic cell closureRef

    ByteArray n bs -> do
        -- allocate the byte array
        arrLength <- handleFromSint64 ctx (fromIntegral (B.length bs)) 64
        arrType <- getID ctx muString
        arrRef <- newHybrid ctx arrType arrLength

        -- populate the array
        ptr <- pin ctx arrRef >>= handleToPtr ctx 
        liftIO $ B.useAsCStringLen bs $ \(strptr, strlen) -> do
            F.copyBytes (F.castPtr ptr :: C.CString) strptr strlen
        unpin ctx arrRef

        -- attach it to the global cell
        cell <- handleFromGlobal ctx =<< getID ctx n
        store ctx muOrdNotAtomic cell arrRef
286 287


288
-------------------------------------------------- * Mu interface
289

290 291 292 293 294 295 296
-- | Get a unique number from the Mu monad.
nextMuUnique :: Mu Int
nextMuUnique = do
    currentID <- gets _uniqueID
    uniqueID %= (+1)
    return currentID

297 298
-- | Emit a closure with a payload and an entry function.  The closure itself
-- will be allocated during build time.
299 300
closure :: (Uniquable a, NamedThing a)
        => a -> [UnboxedData] -> Mu FunctionName -> Mu ()
301
closure name payload code = do
302
    func_name <- code
303 304 305 306
    topLevels %= (|> Closure name' payload func_name)
    definitions %= (|> GlobalCell name' muClosureRef)
  where
    name' = closureNameOf name
307 308


309 310 311 312
-- | Emit a string.  The string itself will be allocated during build time.
string :: GHC.Id -> B.ByteString -> Mu GlobalCellName
string name bytes = do
    name' <- do
313 314
        uid <- nextMuUnique
        return $ fromString $ stringify name <> "_str" <> show uid
315 316 317 318 319 320
    definitions %= (|> GlobalCell name' muStringRef)
    topLevels %= (|> ByteArray name' bytes)
    return name'


-- | Emit a type definition.
321 322 323 324
typedef :: TypedefName -> Type -> Mu ()
typedef n ty = definitions %= (|> TypeDefinition n ty)


325
-- | Emit a constant definition.
326 327 328 329
constant :: ConstantName -> TypedefName -> ConstConstructor -> Mu ()
constant n ty ctor = definitions %= (|> Constant n ty ctor)


330
-- | Emit a function signature definition.
331 332
funcsig :: SignatureName -> [TypedefName] -> [TypedefName] -> Mu ()
funcsig n argtys rettys = definitions %= (|> SignatureDefinition n argtys rettys)
333

334

335
-- | Emit a function (version) definition.
336 337 338 339
funcdef :: FunctionName -> Version -> SignatureName -> Blocks BasicBlock -> Mu ()
funcdef n v sig body = do
    (entry, bblocks) <- runBlocks body
    definitions %= (|> FunctionDefinition n v sig entry (toList bblocks))
340

341

342 343
primfuncdef :: FunctionName -> Blocks BasicBlock -> Mu ()
primfuncdef n body = funcdef n (Version "1") muClosureFunctionSig body
344

345 346 347 348 349 350 351

-------------------------------------------------- * Naming interface

stringify :: (Uniquable a, NamedThing a) => a -> String
stringify a
    | isSystemName name = '@':(stableName name <> "_" <> uniquePart a)
    | otherwise = '@':(stableName name)
352
  where
353 354
    name = getName a
    stableName = zString . zEncodeFS . fsLit . nameStableString
355 356
    uniquePart = show . getUnique

357 358 359
dataConEntry :: GHC.DataCon -> FunctionName
dataConEntry n = FunctionName $ Name $ stringify n <> "_static_entry"

360
closureNameOf :: (Uniquable a, NamedThing a) => a -> GlobalCellName
361
closureNameOf n = GlobalCellName $ Name $ stringify n <> "_closure"
362 363

entryNameOf :: GHC.Id -> FunctionName
364
entryNameOf n = FunctionName $ Name $ stringify n <> "_entry"
365 366 367 368 369 370 371 372


-------------------------------------------------- * Type names

-- | Load a bunch of useful types into the current bundle.  The order of the
-- definitions is arbitrary.
loadTypes :: Mu ()
loadTypes = do
373
    funcsig muClosureFunctionSig [closureTR] [closureTR]
374
    typedef haskellData TagRef64
375
    typedef closureTR TagRef64
376 377 378 379 380
    typedef i64 $ MuInt 64
    constant i64_1 i64 (IntCtor 1)
    typedef i8 $ MuInt 8
    typedef muString $ Hybrid [] i8
    typedef muStringRef $ Ref muString
381
    typedef muClosureMetadata $ MuInt 8
382
    typedef muClosureFunction $ FuncRef muClosureFunctionSig
383 384 385
    typedef muClosure $
        Hybrid [muClosureFunction, muClosureMetadata, muClosureMetadata]
               haskellData
386 387 388 389 390 391 392 393 394 395 396 397 398 399 400
    typedef muClosureRef $ Ref muClosure

i64 :: TypedefName
i64 = "@i64"
i64_1 :: ConstantName
i64_1 = "@i64_1"
i8 :: TypedefName
i8 = "@i8"
muString :: TypedefName
muString = "@string_t"
muStringRef :: TypedefName
muStringRef = "@string_t_ref"

-------------------------------------------------- * STG Closures in Mu

401 402 403 404 405 406 407 408 409 410
-- | A tagged reference to a closure.
closureTR :: TypedefName
closureTR = "@_closure_tagref"

-- | A short field used to store things like the number of elements and the
-- partial-application status of the function. 
muClosureMetadata :: TypedefName
muClosureMetadata = "@_closure_metadata"

-- | The Mu closure itself.
411 412
muClosure :: TypedefName
muClosure = "@_closure"
413

414 415 416 417 418 419 420 421 422
muClosureRef :: TypedefName
muClosureRef = "@_closure_ref"
muClosureFunctionSig :: SignatureName
muClosureFunctionSig = "@_closure_function_sig"
muClosureFunction :: TypedefName
muClosureFunction = "@_closure_function"
haskellData :: TypedefName
haskellData = "@_haskell_data"

423 424 425
-- | An instruction to extract the a particular field from a closure.
getClosureField :: VarName -> ClosureField -> Expression
getClosureField n field = GetFieldIRef False muClosure (fromEnum field) n
426 427 428 429

-- | An instruction to to get the data part from a closure.
getClosureData :: VarName -> Expression
getClosureData = GetVarPartIRef False muClosure
430 431 432 433 434 435 436 437 438 439

-- | An enumeration which corresponds to the position of the fields of Mu
-- closures.
data ClosureField = EntryField
                  | TagField
                  | NArgsField
                  deriving Enum

fieldIx :: ClosureField -> C.CInt
fieldIx = fromIntegral . fromEnum