<commit_before>module Generator(generate) where

import Datatypes
import Text.Printf
import Data.Maybe
import Control.Monad.State
import qualified Data.Map as Map


data CodeGenState =
  CGS {
    availableRegisters :: [Int],
    variables :: (Map.Map String Int),
    generatedCode :: [IR],
    currentLineNumber :: Int
  }
  deriving (Show)


generate :: Program -> [IR]
generate (Program statements) =
  let CGS registers variables code _ = execState
        (mapM generateStatement statements)
        (CGS [7..] Map.empty [] 0)
  in reverse code


generateStatement :: Statement -> State CodeGenState ()
generateStatement (AssignmentStatement (Assignment item _ e1)) = do
  reg1 <- generateExpression e1
  lefthand <- registerForItem item

  emitInstruction $ ThreeIR Plus lefthand (R 0) reg1 False


generateStatement (BuiltinStatement builtinStatement) = do
  let instruction =
        case builtinStatement of
          LoadStatement -> LoadIR
          storeStatement -> StoreIR

  emitInstruction instruction


generateExpression :: Expression -> State CodeGenState IRItem
generateExpression (BinaryExpression Power e1 (ExpressionItem (Immediate power)))
  | power < 1 = error $ "Expected positive exponent (found " ++ show power ++ ")"
  | otherwise = do
    reg <- generateExpression e1
    mapM
      emitInstruction
      $ replicate (power-1) (ThreeIR Multiply reg reg reg False)

    return reg


generateExpression (BinaryExpression op e1 e2) = do
  reg1 <- generateExpression e1
  reg2 <- generateExpression e2

  targetReg <- getRegister
  emitInstruction $ ThreeIR op targetReg reg1 reg2 False

  return targetReg


generateExpression (TernaryExpression e1 e2 e3) = do
  reg1 <- generateExpression e1
  reg2 <- generateExpression e2
  reg3 <- generateExpression e3

  targetReg <- getRegister
  emitInstruction $ ThreeIR Plus (R 6) (R 0) reg1 False
  emitInstruction $ ThreeIR Plus targetReg (R 0) reg2 False
  emitInstruction $ ThreeIR Plus targetReg (R 0) reg3 True

  return targetReg


generateExpression (ExpressionItem item) = do
  CGS (register:registers) variables generatedCode cln <- get
  case item of
    Variable name ->
      case Map.lookup name variables of
        Just val -> return $ R val
        _ -> error $ "Undefined variable '" ++ name  ++ "' around line " ++ show cln

    Register name -> return $ fromRegister name
    Immediate int -> do
      reg <- getRegister
      emitInstruction $ TwoIR reg (I int) False
      return reg

    Constant constant -> do
      reg <- getRegister
      emitInstruction $ TwoIR reg (C constant) False
      return reg

-----------------------------------------------
-- Move the functions below to separate file --
-----------------------------------------------

getRegister :: State CodeGenState IRItem
getRegister = do
  state popRegister
  where popRegister (CGS (r:rs) v c cln) = ((R r), CGS rs v c cln)


registerForItem :: Item -> State CodeGenState IRItem<commit_msg>Remove "Move to separate file" instruction<commit_after>module Generator(generate) where

import Datatypes
import Text.Printf
import Data.Maybe
import Control.Monad.State
import qualified Data.Map as Map


data CodeGenState =
  CGS {
    availableRegisters :: [Int],
    variables :: (Map.Map String Int),
    generatedCode :: [IR],
    currentLineNumber :: Int
  }
  deriving (Show)


generate :: Program -> [IR]
generate (Program statements) =
  let CGS registers variables code _ = execState
        (mapM generateStatement statements)
        (CGS [7..] Map.empty [] 0)
  in reverse code


generateStatement :: Statement -> State CodeGenState ()
generateStatement (AssignmentStatement (Assignment item _ e1)) = do
  reg1 <- generateExpression e1
  lefthand <- registerForItem item

  emitInstruction $ ThreeIR Plus lefthand (R 0) reg1 False


generateStatement (BuiltinStatement builtinStatement) = do
  let instruction =
        case builtinStatement of
          LoadStatement -> LoadIR
          storeStatement -> StoreIR

  emitInstruction instruction


generateExpression :: Expression -> State CodeGenState IRItem
generateExpression (BinaryExpression Power e1 (ExpressionItem (Immediate power)))
  | power < 1 = error $ "Expected positive exponent (found " ++ show power ++ ")"
  | otherwise = do
    reg <- generateExpression e1
    mapM
      emitInstruction
      $ replicate (power-1) (ThreeIR Multiply reg reg reg False)

    return reg


generateExpression (BinaryExpression op e1 e2) = do
  reg1 <- generateExpression e1
  reg2 <- generateExpression e2

  targetReg <- getRegister
  emitInstruction $ ThreeIR op targetReg reg1 reg2 False

  return targetReg


generateExpression (TernaryExpression e1 e2 e3) = do
  reg1 <- generateExpression e1
  reg2 <- generateExpression e2
  reg3 <- generateExpression e3

  targetReg <- getRegister
  emitInstruction $ ThreeIR Plus (R 6) (R 0) reg1 False
  emitInstruction $ ThreeIR Plus targetReg (R 0) reg2 False
  emitInstruction $ ThreeIR Plus targetReg (R 0) reg3 True

  return targetReg


generateExpression (ExpressionItem item) = do
  CGS (register:registers) variables generatedCode cln <- get
  case item of
    Variable name ->
      case Map.lookup name variables of
        Just val -> return $ R val
        _ -> error $ "Undefined variable '" ++ name  ++ "' around line " ++ show cln

    Register name -> return $ fromRegister name
    Immediate int -> do
      reg <- getRegister
      emitInstruction $ TwoIR reg (I int) False
      return reg

    Constant constant -> do
      reg <- getRegister
      emitInstruction $ TwoIR reg (C constant) False
      return reg


getRegister :: State CodeGenState IRItem
getRegister = do
  state popRegister
  where popRegister (CGS (r:rs) v c cln) = ((R r), CGS rs v c cln)


registerForItem :: Item -> State CodeGenState IRItem
{'fraction_non_alphanumeric': 0.08963971729012239, 'fraction_numerical': 0.012066885019824168, 'mean_word_length': 3.1061571125265393, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3900060', 'n_tokens_mistral': 1612, 'n_tokens_neox': 1554, 'n_words': 587}