wip

src/LowToLLVM.hs

@@ -0,0 +1,270 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE NoFieldSelectors #-}
+
+module LowToLLVM where
+
+import Data.ByteString.Builder (Builder)
+import qualified Data.ByteString.Builder as Builder
+import qualified Data.ByteString.Lazy as Lazy
+import Data.HashMap.Lazy (HashMap)
+import qualified Data.HashMap.Lazy as HashMap
+import Data.HashSet (HashSet)
+import qualified Data.HashSet as HashSet
+import Data.Text.Encoding (encodeUtf8Builder)
+import qualified Index.Seq
+import qualified Index.Seq as Index (Seq)
+import qualified Literal
+import Low.PassBy (PassBy)
+import qualified Low.PassBy as PassBy
+import qualified Low.Representation as Representation
+import qualified Low.Syntax as Syntax
+import Name (Name)
+import qualified Name
+import Protolude hiding (IntMap, cast, local, moduleName, repr)
+
+newtype Var = Var Text
+  deriving (Eq, Ord, Show, Hashable)
+
+nameBuilder :: Var -> Builder
+nameBuilder (Var n) = encodeUtf8Builder n
+
+varName :: Var -> Builder
+varName n = "%" <> nameBuilder n
+
+liftedName :: Name.Lifted -> Builder
+liftedName = \case
+  Name.Lifted (Name.Qualified (Name.Module moduleName) (Name.Name name_)) 0 ->
+    "@" <> encodeUtf8Builder moduleName <> "." <> encodeUtf8Builder name_
+  Name.Lifted (Name.Qualified (Name.Module moduleName) (Name.Name name_)) i ->
+    "@" <> encodeUtf8Builder moduleName <> "." <> encodeUtf8Builder name_ <> "$" <> Builder.intDec i
+
+data Operand
+  = Local !Var
+  | Global !Name.Lifted
+  | Constant !Builder
+  deriving (Show)
+
+operand :: Operand -> Builder
+operand = \case
+  Local v -> varName v
+  Global n -> liftedName n
+  Constant c -> c
+
+separate :: Builder -> [Builder] -> Builder
+separate separator = mconcat . intersperse separator
+
+commaSeparate :: [Builder] -> Builder
+commaSeparate = separate ", "
+
+parens :: [Builder] -> Builder
+parens bs = "(" <> commaSeparate bs <> ")"
+
+braces :: [Builder] -> Builder
+braces bs = "{" <> commaSeparate bs <> "}"
+
+brackets :: [Builder] -> Builder
+brackets bs = "[" <> commaSeparate bs <> "]"
+
+wordBits :: (Num a) => a
+wordBits = 64
+
+wordSizedInt :: Builder
+wordSizedInt = "i" <> Builder.intDec wordBits
+
+type Assembler = State AssemblerState
+
+data AssemblerState = AssemblerState
+  { usedGlobals :: HashSet Name.Lifted
+  , usedLLVMGlobals :: HashMap Text Builder
+  , usedLocals :: HashMap Var Int
+  , instructions :: Builder
+  , basicBlockName :: Var
+  , basicBlocks :: Builder
+  }
+
+runAssembler :: Assembler a -> (a, (HashSet Name.Lifted, HashMap Text Builder))
+runAssembler =
+  second (\s -> (s.usedGlobals, s.usedLLVMGlobals))
+    . flip
+      runState
+      AssemblerState
+        { usedLocals = mempty
+        , usedGlobals = mempty
+        , usedLLVMGlobals = mempty
+        , instructions = mempty
+        , basicBlocks = mempty
+        , basicBlockName = panic "AssemblyToLLVM: not in a basic block"
+        }
+
+llvmType :: PassBy -> Builder
+llvmType = \case
+  PassBy.Reference -> "{ ptr, ptr }"
+  PassBy.Value repr ->
+    "{ ["
+      <> Builder.intDec repr.pointers
+      <> " x "
+      <> wordSizedInt
+      <> "], ["
+      <> Builder.intDec repr.nonPointerBytes
+      <> " x i8] }"
+
+emitInstruction :: Builder -> Assembler ()
+emitInstruction instruction =
+  modify \s -> s {instructions = s.instructions <> "\n  " <> instruction}
+
+startBlock :: Var -> Assembler ()
+startBlock basicBlockName =
+  modify \s -> s {basicBlockName}
+
+endBlock :: Builder -> Assembler ()
+endBlock terminator =
+  modify \s ->
+    s
+      { instructions = mempty
+      , basicBlockName = panic "AssemblyToLLVM: not in a basic block"
+      , basicBlocks =
+          s.basicBlocks
+            <> "\n\n"
+            <> nameBuilder s.basicBlockName
+            <> ":"
+            <> s.instructions
+            <> "\n  "
+            <> terminator
+      }
+
+freshVar :: Name -> Assembler Var
+freshVar (Name.Name nameText) = do
+  usedLocals <- gets (.usedLocals)
+  let (i, usedNames') =
+        HashMap.alterF
+          ( \case
+              Nothing -> (0, Just 1)
+              Just j -> (j, Just $ j + 1)
+          )
+          (Var nameText)
+          usedLocals
+  modify \s -> s {usedLocals = usedNames'}
+  pure $ Var if i == 0 then nameText else nameText <> "$" <> (show i :: Text)
+
+declareGlobal :: Name.Lifted -> Assembler ()
+declareGlobal name =
+  modify \s ->
+    s {usedGlobals = HashSet.insert name s.usedGlobals}
+
+declareLLVMGlobal :: Text -> Builder -> Assembler ()
+declareLLVMGlobal name decl =
+  modify \s ->
+    s {usedLLVMGlobals = HashMap.insert name decl s.usedLLVMGlobals}
+
+-------------------------------------------------------------------------------
+
+assembleModule :: [(Name.Lifted, Syntax.Definition)] -> Lazy.ByteString
+assembleModule definitions = do
+  let (assembledDefinitions, allUsedGlobals) =
+        unzip $ foreach definitions $ runAssembler . uncurry assembleDefinition
+      (usedGlobals, usedLLVMGlobals) = bimap mconcat mconcat $ unzip allUsedGlobals
+      assembledDefinitions' = concat assembledDefinitions
+  Builder.toLazyByteString $
+    separate "\n\n" $
+      HashMap.elems usedLLVMGlobals <> map snd assembledDefinitions'
+
+type Environment v = Index.Seq v (PassBy, Operand)
+
+assembleDefinition :: Name.Lifted -> Syntax.Definition -> Assembler [(Name.Lifted, Builder)]
+assembleDefinition name definition =
+  case definition of
+    Syntax.FunctionDefinition function ->
+      pure <$> assembleFunction name Index.Seq.Empty function
+    Syntax.ConstantDefinition _ _ -> undefined
+
+assembleFunction
+  :: Name.Lifted
+  -> Environment v
+  -> Syntax.Function v
+  -> Assembler (Name.Lifted, Builder)
+assembleFunction functionName env = \case
+  Syntax.Parameter name passBy function -> do
+    var <- freshVar name
+    assembleFunction functionName (env Index.Seq.:> (passBy, Local var)) function
+  Syntax.Body passReturnBy term -> do
+    let parameters = second fromLocal <$> Index.Seq.toSeq env
+    (result, restore) <- assembleTerm env term
+    emitInstruction $ "ret " <> llvmType passReturnBy <> " " <> operand result
+    basicBlocks <- gets (.basicBlocks)
+    pure
+      ( functionName
+      , "define "
+          <> linkage
+          <> "fastcc "
+          <> llvmType passReturnBy
+          <> " "
+          <> liftedName functionName
+          <> parens
+            [ llvmType passBy <> " " <> varName parameter
+            | (passBy, parameter) <- toList parameters
+            ]
+          <> " align "
+          <> Builder.intDec alignment
+          <> " {"
+          <> basicBlocks
+          <> "\n}"
+      )
+  where
+    fromLocal (Local l) = l
+    fromLocal _ = panic "non-local function parameter"
+    linkage =
+      case functionName of
+        Name.Lifted _ 0 ->
+          ""
+        _ ->
+          "private "
+    alignment :: (Num a) => a
+    alignment = 8
+
+assembleTerm :: Environment v -> Syntax.Term v -> Maybe Name -> Assembler (Operand, Maybe Var)
+assembleTerm env nameSuggestion = \case
+  Syntax.Operand operand -> (,Nothing) <$> assembleOperand env operand
+  Syntax.Let passBy name term body -> do
+    (termResult, termStack) <- assembleTerm env (Just name) term
+    (bodyResult, bodyStack) <- assembleTerm (env Index.Seq.:> (passBy, termResult)) body
+    mapM_ restoreStack termStack
+    mapM_ restoreStack bodyStack
+    pure (result, Nothing)
+  Syntax.Seq term1 term2 -> do
+    (_, stack1) <- assembleTerm env Nothing term1
+    (result, stack2) <- assembleTerm env nameSuggestion term2
+    mapM_ restoreStack stack1
+    mapM_ restoreStack stack2
+    pure (result, Nothing)
+  Syntax.Case {} -> _
+  Syntax.Call name args -> _
+  Syntax.StackAllocate operand -> _
+  Syntax.HeapAllocate con size -> _
+  Syntax.HeapPayload pointer -> _
+  Syntax.PointerTag pointer -> _
+  Syntax.Offset base size -> _
+  Syntax.Copy dst src size -> _
+  Syntax.Store dst src repr -> _
+  Syntax.Load src repr -> _
+
+assembleOperand :: Environment v -> Syntax.Operand v -> Assembler Operand
+assembleOperand env = \case
+  Syntax.Var index -> _
+  Syntax.Global global -> _
+  Syntax.Literal literal -> _
+  Syntax.Representation repr -> _
+  Syntax.Tag tag -> _
+  Syntax.Undefined repr -> _
+
+saveStack :: Assembler Var
+saveStack = undefined
+
+restoreStack :: Var -> Assembler ()
+restoreStack = undefined