content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Glean.Bytecode.Generate.Haskell (main)
  where
  
  import Data.Char (toLower)
  import Data.List (intercalate, intersperse, stripPrefix)
  import Data.Maybe
  import Data.Text (Text)
  import qualified Data.Text as Text
  import qualified Data.Text.IO as Text
  import System.Directory
  import System.Environment
  import System.Exit (die)
  import System.FilePath
  
  import Glean.Bytecode.Generate.Instruction
  import Glean.Bytecode.Types
  
  -- We generate two modules:
  --
  -- Instruction: contains the Insn type and some utility functions
  --
  -- Issue: a typed interface for generating instructions based on the MonadInsn
  -- class from Glean.RTS.Bytecode.MonadInsn.
  
  list :: Text -> Text -> [Text] -> [Text]
  list first others (x:xs) = first <> x : map (others <>) xs
  list _ _ [] = []
  
  modPfx :: Text
  modPfx = "Glean.RTS.Bytecode.Gen"
  
  modPath :: FilePath
  modPath = Text.unpack $ Text.replace "." "/" modPfx
  
  main :: IO ()
  main = do
    args <- getArgs
    dir <- case args of
      [arg] | Just dir <- stripPrefix "--install_dir=" arg ->
        return $ dir </> modPath
      _ -> die "invalid arguments"
    createDirectoryIfMissing True dir
    genModule dir "Instruction"
      []
      [ "Insn(..)"
      , "mapLabels"
      , "insnLabels"
      , "insnSize"
      , "insnWords"
      , "insnShow" ] $
      intercalate [""]
      [ [ "import Data.Word (Word64)"
        , "import Text.Show (showListWith)"
        , "import qualified Glean.Bytecode.Decode as D"
        , "import Glean.Bytecode.Types" ]
      , genInsnType
      , genMapLabels
      , genInsnLabels
      , genInsnSize
      , genInsnWords
      , genDecodable
      , genInsnShow ]
  
    genModule dir "Issue" [] (map (varName . insnName) instructions) $
      [ "import Data.ByteString (ByteString)"
      , "import Data.Word (Word64)"
      , "import Glean.RTS.Bytecode.Gen.Instruction (Insn(..))"
      , "import Glean.RTS.Bytecode.Code"
      , "import Glean.Bytecode.Types"
      , "" ]
      ++ genIssue
  
    genModule dir "Version" [] ["version", "lowestSupportedVersion"]
      [ "version :: Int"
      , "version = " <> Text.pack (show version)
      , ""
      , "lowestSupportedVersion :: Int"
      , "lowestSupportedVersion = " <> Text.pack (show lowestSupportedVersion)
      ]
  
  genModule :: FilePath -> String -> [Text] -> [Text] -> [Text] -> IO ()
  genModule path name exts es ls =
    Text.writeFile (path </> name <.> "hs")
    $ Text.unlines $
    [ "-- @" <> "generated" ]
    ++
    [ "{-# LANGUAGE " <> Text.intercalate ", " exts <> " #-}" | not (null exts)]
    ++
    [ "{-# OPTIONS_GHC -Wno-unused-matches #-}"
    , "module " <> modPfx <> "." <> Text.pack name ]
    ++ list "  ( " "  , " es ++
    [ "  ) where"
    , "" ]
    ++ ls
  
  -- | Generate the Insn type. We parametrise over the types of registers and
  -- labels.
  genInsnType :: [Text]
  genInsnType = "data Insn where" : map genInsn instructions
    where
      genInsn Insn{..} =
        "  "
          <> insnName
          <> " :: "
          <> Text.concat
              [ty <> " -> " | arg <- insnArgs, ty <- genArgTy $ argTy arg]
          <> "Insn"
  
      genArgTy (Imm ImmOffset) = ["{-# UNPACK #-} !Label"]
      genArgTy (Imm ImmLit) = ["{-# UNPACK #-} !Literal"]
      genArgTy (Imm ImmWord) = ["{-# UNPACK #-} !Word64"]
      genArgTy (Reg var ty _) = ["{-# UNPACK #-} !(" <> showRegTy var ty <> ")"]
      genArgTy Offsets = ["[Label]"]
      genArgTy (Regs tys) = ["(Register " <> showTy ty <> ")" | ty <- tys]
  
  showTy :: Ty -> Text
  showTy ty = case ty of
    Fun args ->
      "('Fun '[ " <> Text.intercalate "," (map showTy args) <> " ])"
    _ -> "'" <> Text.pack (show ty)
  
  showRegTy :: Maybe Text -> Ty -> Text
  showRegTy var ty = "Register " <> fromMaybe (showTy ty) var
  
  genMapLabels :: [Text]
  genMapLabels =
    "mapLabels :: (Label -> Label) -> Insn -> Insn" :
    [ "mapLabels f ("
        <> Text.unwords (insnName insn : map argName (insnArgs insn))
        <> ") = "
        <> Text.unwords (insnName insn : exprs)
      | insn <- instructions
      , let exprs = map mkMap (insnArgs insn)
      , any (\(expr,arg) -> expr /= argName arg) $ zip exprs $ insnArgs insn ]
    ++
    [ "mapLabels _ insn = insn"]
    where
      mkMap (Arg name (Imm ImmOffset)) = "(f " <> name <> ")"
      mkMap (Arg name Offsets) = "(map f " <> name <> ")"
      mkMap (Arg name _) = name
  
  genInsnLabels :: [Text]
  genInsnLabels =
    "insnLabels :: Insn -> [Label]" :
    [ "insnLabels ("
        <> Text.unwords (insnName insn : map argName (insnArgs insn))
        <> ") = "
        <> Text.intercalate " ++ " exprs
      | insn <- instructions
      , let exprs = mapMaybe mkList (insnArgs insn)
      , not $ null exprs ]
    ++
    [ "insnLabels _ = []"]
    where
      mkList (Arg name (Imm ImmOffset)) = Just $ "[" <> name <> "]"
      mkList (Arg name Offsets) = Just name
      mkList Arg{} = Nothing
  
  -- | Generates a function that yields the size of an instruction in words.
  genInsnSize :: [Text]
  genInsnSize =
    "insnSize :: Insn -> Word64" :
    [ "insnSize "
        <> mkPat insnName (concatMap mkArg insnArgs)
        <> " = "
        <> showSize (foldr argSize (1::Int,"") insnArgs)
      | Insn{..} <- instructions ]
    where
      mkPat name [] = name
      mkPat name args
        | all (=="_") args = name <> "{}"
        | otherwise = "(" <> Text.unwords (name : args) <> ")"
  
      mkArg (Arg name Offsets) = [name]
      mkArg (Arg _ (Regs tys)) = replicate (length tys) "_"
      mkArg _ = ["_"]
  
      showSize (c,v) = Text.pack (show c) <> v
  
      argSize (Arg name Offsets) (c,v) =
        (c+1, v <> " + fromIntegral (length " <> name <> ")")
      argSize (Arg _ (Regs tys)) (c,v) = (c + length tys, v)
      argSize _ (c,v) = (c+1, v)
  
  -- | Generates a function that encodes an instruction as a list of words.
  genInsnWords :: [Text]
  genInsnWords =
     [ "insnWords"
     , "  :: (forall ty. Register ty -> Word64) -> (Label -> Word64) -> Insn -> [Word64]" ]
     ++
     [ "insnWords fromReg fromLabel ("
        <> insnPattern id insn
        <> ") = ["
        <> Text.pack (show op)
        <> Text.concat (map argWords (insnArgs insn))
        <> "]" | (op, insn) <- zip [0 :: Int ..] instructions ]
      where
        argWords (Arg name (Imm ImmOffset)) = ", fromLabel " <> name
        argWords (Arg name (Imm ImmLit)) = ", fromLiteral " <> name
        argWords (Arg name (Imm ImmWord)) = ", " <> name
        argWords (Arg name Reg{}) = ", fromReg " <> name
        argWords (Arg name Offsets) =
          ", fromIntegral (length " <> name <> ")] ++ map fromLabel " <> name
            <> " ++ ["
        argWords (Arg name (Regs tys)) =
          Text.concat [", fromReg " <> reg | reg <- regNames name tys]
  
  insnPattern :: (Text -> Text) -> Insn -> Text
  insnPattern cname Insn{..} =
    Text.unwords $ cname insnName : concatMap argNames insnArgs
    where
      argNames (Arg name (Regs tys)) = regNames name tys
      argNames (Arg name _) = [name]
  
  regNames :: Text -> [Ty] -> [Text]
  regNames name tys = [name <> Text.pack (show i) | i <- [1 .. length tys]]
  
  varName :: Text -> Text
  varName name
    | Just (c,rest) <- Text.uncons name = Text.cons (toLower c) rest
    | otherwise = name
  
  genIssue :: [Text]
  genIssue = intercalate [""]
    [ [ varName insnName
          <> " :: "
          <> context insnContext
          <> Text.unwords
              (intersperse "->"
                $ concatMap (genArgType . argTy) insnArgs ++ ["Code ()"])
      , insnPattern varName insn <> " = do"
      ] ++ map ("  " <>)
        (mapMaybe literal insnArgs ++
        [ issue insnEffects
            <> " $ " <> Text.unwords (insnName : concatMap genArgRef insnArgs) ])
      | insn@Insn{..} <- instructions ]
    where
      context [] = ""
      context [c] = c <> " => "
      context cs = "(" <> Text.intercalate ", " cs <> ") => "
  
      genArgType (Imm ImmLit) = ["ByteString"]
      genArgType (Imm ImmOffset) = ["Label"]
      genArgType (Imm ImmWord) = ["Word64"]
      genArgType (Reg var ty _) = [showRegTy var ty]
      genArgType Offsets = ["[Label]"]
      genArgType (Regs tys) = ["Register " <> showTy ty | ty <- tys]
  
      genArgRef (Arg name (Imm ImmLit)) = [name <> "_i"]
      genArgRef (Arg name (Regs tys)) = regNames name tys
      genArgRef (Arg name _) = [name]
  
      literal (Arg name (Imm ImmLit)) = Just $
        name <> "_i <- literal " <> name
      literal _ = Nothing
  
      issue effects
        | EndBlock `elem` effects = "issueEndBlock"
        | otherwise = "issue"
  
  genDecodable :: [Text]
  genDecodable =
    [ "instance D.Decodable Insn where"
    , "  decode = do"
    , "    op <- D.decode"
    , "    case (op :: Word64) of" ]
    ++
    [ Text.concat $ [ "      ", Text.pack (show i), " -> " ] ++
      case sum (map arity insnArgs) of
        0 -> ["pure ", insnName]
        n -> insnName : " <$> D.decode" : replicate (n-1) " <*> D.decode"
      | (i, Insn{..}) <- zip [0 :: Int ..] instructions ]
    ++
    [ "      _ -> fail $ \"invalid opcode \" ++ show op"]
    where
      arity (Arg _ (Regs tys)) = length tys
      arity _ = 1
  
  genInsnShow :: [Text]
  genInsnShow =
    [ "insnShow"
    , "  :: [String]"
    , "  -> (Label -> String)"
    , "  -> (forall t. Register t -> String)"
    , "  -> Insn"
    , "  -> String" ]
    ++
    [ "insnShow syscalls showLabel showReg ("
        <> insnPattern id insn
        <> ") = concat [" <> name
        <> Text.intercalate ", \",\""
            [", \' \' : " <> w | arg <- args, w <- showArg arg]
        <> "]" | insn@Insn{..} <- instructions
               , let (name, args) = mkNameArgs insnName insnArgs ]
      where
        showArg (Arg name (Imm ImmOffset)) = ["showLabel " <> name]
        showArg (Arg name (Imm ImmLit)) = ["show (fromLiteral " <> name <> ")"]
        showArg (Arg name (Imm ImmWord)) = ["show " <> name]
        showArg (Arg name Reg{}) = ["showReg " <> name]
        showArg (Arg name Offsets) =
          ["showListWith (showString . showLabel) " <> name <> " \"\""]
        showArg (Arg name (Regs tys)) =
          ["showReg " <> reg | reg <- regNames name tys]
        mkNameArgs name args
          | "CallFun" `Text.isPrefixOf` name =
            ("(syscalls !! fromIntegral (fromRegister fun))", tail args)
          | otherwise = ("\"" <> name <> "\"", args)
  