Add IR support for Wasm globals (#727)

asterius/src/Asterius/Backends/Binaryen.hs

@@ -94,6 +94,15 @@ marshalReturnTypes vts = case vts of
       "binaryen doesn't support multi-value yet: failed to marshal "
         <> show vts
 
+marshalMutability :: Mutability -> Int8
+marshalMutability = \case
+  Immutable -> 0
+  Mutable -> 1
+
+marshalGlobalType :: GlobalType -> (Binaryen.Type, Int8)
+marshalGlobalType GlobalType {..} =
+  (marshalValueType globalValueType, marshalMutability globalMutability)
+
 marshalUnaryOp :: UnaryOp -> Binaryen.Op
 marshalUnaryOp op = case op of
   ClzInt32 -> Binaryen.clzInt32
@@ -356,6 +365,19 @@ marshalExpression e = case e of
     v <- marshalExpression value
     m <- askModuleRef
     lift $ Binaryen.localTee m (coerce index) v $ marshalValueType valueType
+  GetGlobal {..} -> do
+    m <- askModuleRef
+    a <- askArena
+    lift $ do
+      gbl <- marshalBS a (entityName globalSymbol)
+      Binaryen.globalGet m gbl $ marshalValueType valueType
+  SetGlobal {..} -> do
+    v <- marshalExpression value
+    m <- askModuleRef
+    a <- askArena
+    lift $ do
+      gbl <- marshalBS a (entityName globalSymbol)
+      Binaryen.globalSet m gbl v
   Load {..} -> do
     p <- marshalExpression ptr
     m <- askModuleRef
@@ -596,6 +618,36 @@ marshalMemoryImport m MemoryImport {..} = do
     ebp <- marshalBS a externalBaseName
     Binaryen.addMemoryImport m inp emp ebp 0
 
+marshalGlobalImport :: Binaryen.Module -> GlobalImport -> CodeGen ()
+marshalGlobalImport m GlobalImport {..} = do
+  a <- askArena
+  lift $ do
+    inp <- marshalBS a internalName
+    emp <- marshalBS a externalModuleName
+    ebp <- marshalBS a externalBaseName
+    let (ty, mut) = marshalGlobalType globalType
+    Binaryen.addGlobalImport m inp emp ebp ty (fromIntegral mut)
+
+marshalGlobalExport ::
+  Binaryen.Module -> GlobalExport -> CodeGen Binaryen.Export
+marshalGlobalExport m GlobalExport {..} = do
+  a <- askArena
+  lift $ do
+    inp <- marshalBS a internalName
+    enp <- marshalBS a externalName
+    Binaryen.addGlobalExport m inp enp
+
+marshalGlobal ::
+  BS.ByteString -> Global -> CodeGen Binaryen.Global
+marshalGlobal k Global {..} = do
+  let (ty, mut) = marshalGlobalType globalType
+  e <- marshalExpression globalInit
+  m <- askModuleRef
+  a <- askArena
+  lift $ do
+    ptr <- marshalBS a k
+    Binaryen.addGlobal m ptr ty mut e
+
 marshalModule :: Bool -> SM.SymbolMap Int64 -> Module -> IO Binaryen.Module
 marshalModule tail_calls sym_map hs_mod@Module {..} = do
   m <- Binaryen.Module.create
@@ -621,6 +673,10 @@ marshalModule tail_calls sym_map hs_mod@Module {..} = do
       forM_ functionImports $ \fi@FunctionImport {..} ->
         marshalFunctionImport m (ftps M.! functionType) fi
       forM_ functionExports $ marshalFunctionExport m
+      forM_ (SM.toList globalMap) $
+        \(k, global) -> marshalGlobal (entityName k) global
+      forM_ globalImports $ marshalGlobalImport m
+      forM_ globalExports $ marshalGlobalExport m
       marshalFunctionTable m tableSlots functionTable
       marshalTableImport m tableImport
       marshalMemorySegments memoryMBlocks memorySegments

asterius/src/Asterius/Backends/WasmToolkit.hs

@@ -12,6 +12,7 @@
 
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE StrictData #-}
@@ -48,7 +49,7 @@ import qualified Language.WebAssembly.WireFormat as Wasm
 
 data MarshalError
   = DuplicateFunctionImport
-  | DuplicateGlobalImport             -- ^ Currently unused.
+  | DuplicateGlobalImport
   | InvalidParameterType              -- ^ Currently unused.
   | InvalidLocalType                  -- ^ Currently unused.
   | UnsupportedExpression Expression
@@ -59,7 +60,8 @@ instance Exception MarshalError
 data ModuleSymbolTable
   = ModuleSymbolTable
       { functionTypeSymbols :: Map.Map FunctionType Wasm.FunctionTypeIndex,
-        functionSymbols :: Map.Map BS.ByteString Wasm.FunctionIndex
+        functionSymbols :: Map.Map BS.ByteString Wasm.FunctionIndex,
+        globalSymbols :: Map.Map BS.ByteString Wasm.GlobalIndex
       }
 
 makeModuleSymbolTable ::
@@ -71,20 +73,32 @@ makeModuleSymbolTable m@Module {..} = do
       _func_syms = Map.keys functionMap'
       _func_conflict_syms = _func_import_syms `intersect` _func_syms
       _func_types = generateWasmFunctionTypeSet m
-  if _has_dup _func_import_syms
-    then throwError DuplicateFunctionImport
-    else pure ModuleSymbolTable
-      { functionTypeSymbols =
-          Map.fromDistinctAscList $
-            zip
-              (Set.toList _func_types)
-              (coerce [0 :: Word32 ..]),
-        functionSymbols =
-          Map.fromList $
-            zip
-              (_func_import_syms <> _func_syms)
-              (coerce [0 :: Word32 ..])
-      }
+      _gbl_import_syms =
+        [internalName | GlobalImport {..} <- globalImports]
+      _gbl_syms = map entityName $ SM.keys globalMap
+  if  | _has_dup _func_import_syms ->
+        throwError DuplicateFunctionImport
+      | _has_dup _gbl_import_syms ->
+        throwError DuplicateGlobalImport
+      | otherwise ->
+        pure
+          ModuleSymbolTable
+            { functionTypeSymbols =
+                Map.fromDistinctAscList $
+                  zip
+                    (Set.toList _func_types)
+                    (coerce [0 :: Word32 ..]),
+              functionSymbols =
+                Map.fromList $
+                  zip
+                    (_func_import_syms <> _func_syms)
+                    (coerce [0 :: Word32 ..]),
+              globalSymbols =
+                Map.fromList $
+                  zip
+                    (_gbl_import_syms <> _gbl_syms)
+                    (coerce [0 :: Word32 ..])
+            }
 
 makeValueType :: ValueType -> Wasm.ValueType
 makeValueType vt = case vt of
@@ -93,6 +107,31 @@ makeValueType vt = case vt of
   F32 -> Wasm.F32
   F64 -> Wasm.F64
 
+makeMutability :: Mutability -> Wasm.Mutability
+makeMutability m = case m of
+  Immutable -> Wasm.Const
+  Mutable -> Wasm.Var
+
+makeGlobalType :: GlobalType -> Wasm.GlobalType
+makeGlobalType GlobalType {..} =
+  Wasm.GlobalType
+    { Wasm.globalValueType = makeValueType globalValueType,
+      Wasm.globalMutability = makeMutability globalMutability
+    }
+
+makeGlobal ::
+  (MonadError MarshalError m, MonadReader MarshalEnv m) =>
+  Global ->
+  m Wasm.Global
+makeGlobal Global {..} = do
+  let ty = makeGlobalType globalType
+  e <- Wasm.Expression . bagToList <$> makeInstructions globalInit
+  pure
+    Wasm.Global
+      { globalType = ty,
+        globalInitialValue = e
+      }
+
 makeTypeSection ::
   MonadError MarshalError m => Module -> ModuleSymbolTable -> m Wasm.Section
 makeTypeSection Module {} ModuleSymbolTable {..} = do
@@ -142,6 +181,14 @@ makeImportSection Module {..} ModuleSymbolTable {..} = pure Wasm.ImportSection
               }
             | FunctionImport {..} <- functionImports
           ]
+        ++ [ Wasm.Import
+               { moduleName = coerce $ SBS.toShort externalModuleName,
+                 importName = coerce $ SBS.toShort externalBaseName,
+                 importDescription =
+                   Wasm.ImportGlobal $ makeGlobalType globalType
+               }
+             | GlobalImport {..} <- globalImports
+           ]
   }
 
 makeFunctionSection ::
@@ -153,6 +200,25 @@ makeFunctionSection Module {..} ModuleSymbolTable {..} = pure Wasm.FunctionSecti
       ]
   }
 
+makeGlobalSection ::
+  MonadError MarshalError f =>
+  Bool ->
+  SM.SymbolMap Int64 ->
+  Module ->
+  ModuleSymbolTable ->
+  f Wasm.Section
+makeGlobalSection tail_calls sym_map Module {..} _module_symtable = do
+  let env = MarshalEnv
+        { envAreTailCallsOn = tail_calls,
+          envSymbolMap = sym_map,
+          envModuleSymbolTable = _module_symtable,
+          envDeBruijnContext = emptyDeBruijnContext,
+          envLclContext = emptyLocalContext
+        }
+  fmap Wasm.GlobalSection $
+    flip runReaderT env $
+      mapM makeGlobal (SM.elems globalMap)
+
 makeExportSection ::
   MonadError MarshalError m => Module -> ModuleSymbolTable -> m Wasm.Section
 makeExportSection Module {..} ModuleSymbolTable {..} = pure Wasm.ExportSection
@@ -166,6 +232,15 @@ makeExportSection Module {..} ModuleSymbolTable {..} = pure Wasm.ExportSection
           }
         | FunctionExport {..} <- functionExports
       ]
+        ++ [ Wasm.Export
+               { exportName = coerce $ SBS.toShort externalName,
+                 exportDescription =
+                   Wasm.ExportGlobal $
+                     globalSymbols
+                       Map.! internalName
+               }
+             | GlobalExport {..} <- globalExports
+           ]
   }
 
 makeElementSection ::
@@ -546,6 +621,20 @@ makeInstructions expr =
         v `snocBag` Wasm.TeeLocal
           { teeLocalIndex = idx
           }
+    GetGlobal {..} -> do
+      ModuleSymbolTable {..} <- askModuleSymbolTable
+      pure $ unitBag Wasm.GetGlobal
+        { getGlobalIndex =
+            globalSymbols Map.! entityName globalSymbol
+        }
+    SetGlobal {..} -> do
+      v <- makeInstructions value
+      ModuleSymbolTable {..} <- askModuleSymbolTable
+      pure $
+        v `snocBag` Wasm.SetGlobal
+          { setGlobalIndex =
+              globalSymbols Map.! entityName globalSymbol
+          }
     Load {..} -> do
       let _mem_arg = Wasm.MemoryArgument
             { memoryArgumentAlignment = 0,
@@ -751,6 +840,7 @@ makeModule tail_calls sym_map m = do
   _type_sec <- makeTypeSection m _module_symtable
   _import_sec <- makeImportSection m _module_symtable
   _func_sec <- makeFunctionSection m _module_symtable
+  _gbl_sec <- makeGlobalSection tail_calls sym_map m _module_symtable
   _export_sec <- makeExportSection m _module_symtable
   _elem_sec <- makeElementSection m _module_symtable
   _code_sec <- makeCodeSection tail_calls sym_map m _module_symtable
@@ -760,6 +850,7 @@ makeModule tail_calls sym_map m = do
       [ _type_sec,
         _import_sec,
         _func_sec,
+        _gbl_sec,
         _export_sec,
         _elem_sec,
         _code_sec,

asterius/src/Asterius/Builtins.hs

@@ -8,6 +8,8 @@ module Asterius.Builtins
     rtsAsteriusModule,
     rtsFunctionImports,
     rtsFunctionExports,
+    rtsGlobalImports,
+    rtsGlobalExports,
     emitErrorMessage,
     wasmPageSize,
     generateWrapperFunction,
@@ -681,6 +683,12 @@ rtsFunctionExports debug =
       }
   ]
 
+rtsGlobalImports :: [GlobalImport]
+rtsGlobalImports = mempty
+
+rtsGlobalExports :: [GlobalExport]
+rtsGlobalExports = mempty
+
 emitErrorMessage :: [ValueType] -> BS.ByteString -> Expression
 emitErrorMessage vts ev = Barf {barfMessage = ev, barfReturnTypes = vts}
 

asterius/src/Asterius/Passes/GCSections.hs

@@ -66,6 +66,7 @@ gcModule mod_syms export_funcs m =
   AsteriusModule
     { staticsMap = statics,
       functionMap = functions,
+      globalsMap = globals,
       sptMap = spt_map,
       ffiMarshalState =
         FFIMarshalState
@@ -76,6 +77,7 @@ gcModule mod_syms export_funcs m =
   where
     statics = staticsMap m `SM.restrictKeys` mod_syms
     functions = functionMap m `SM.restrictKeys` mod_syms
+    globals = globalsMap m `SM.restrictKeys` mod_syms
     spt_map = sptMap m `SM.restrictKeys` mod_syms
     -- Since each JSFFI import comes in two parts (a function import and a
     -- wrapper function), we only keep the import whose wrapper function is

asterius/src/Asterius/Resolve.hs

@@ -93,6 +93,9 @@ resolveAsteriusModule debug bundled_ffi_state m_globals_resolved func_start_addr
             externalBaseName = "table"
           },
         tableSlots = table_slots,
+        globalImports = rtsGlobalImports,
+        globalExports = rtsGlobalExports,
+        globalMap = globalsMap m_globals_resolved, -- Copy as-is.
         memorySegments = segs,
         memoryImport = MemoryImport
           { externalModuleName = "WasmMemory",

asterius/src/Asterius/TypeInfer.hs

@@ -27,6 +27,9 @@ infer expr = case expr of
   CallIndirect {} -> [I64]
   GetLocal {..} -> [valueType]
   SetLocal {} -> []
+  TeeLocal {..} -> [valueType]
+  GetGlobal {..} -> [valueType]
+  SetGlobal {} -> []
   Load {..} -> [valueType]
   Store {} -> []
   ConstI32 {} -> [I32]
@@ -41,7 +44,6 @@ infer expr = case expr of
   Symbol {} -> [I64]
   UnresolvedGetLocal {..} -> [typeOfUnresolvedLocalReg unresolvedLocalReg]
   UnresolvedSetLocal {} -> []
-  TeeLocal {..} -> [valueType]
   Drop {} -> []
   -- ReturnCall and ReturnCallIndirect are generated when we exit a cmm
   -- function and jump to the next, and all cmm functions have type: [] -> [].