<commit_before>{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module LLVM.General.Pretty (
  ppllvm,
) where

import GHC.Word

import LLVM.General.Typed

import LLVM.General.AST
import LLVM.General.AST.Global
import LLVM.General.AST.Type

import LLVM.General.AST.Attribute
import qualified LLVM.General.AST.Linkage as L
import qualified LLVM.General.AST.Visibility as V
import qualified LLVM.General.AST.CallingConvention as CC
import qualified LLVM.General.AST.Constant as C
import qualified LLVM.General.AST.FloatingPointPredicate as FP
import qualified LLVM.General.AST.IntegerPredicate as IP
import qualified LLVM.General.AST.AddrSpace as AS
import qualified LLVM.General.AST.Float as F

import Data.String
import Text.Printf
import Text.PrettyPrint

import Data.Char (chr, ord, isControl, isLetter, isDigit)
import Data.List (intersperse)
import Numeric (showHex)

-------------------------------------------------------------------------------
-- Utils
-------------------------------------------------------------------------------

parensIf ::  Bool -> Doc -> Doc
parensIf True = parens
parensIf False = id

commas :: [Doc] -> Doc
commas  = hsep . punctuate (char ',')

colons :: [Doc] -> Doc
colons  = hcat . intersperse (char ':')

hlinecat :: [Doc] -> Doc
hlinecat = vcat . intersperse (text "")

wrapbraces :: Doc -> Doc -> Doc
wrapbraces leadIn x = (leadIn <> char '{') $+$ ( x ) $+$ char '}'

local :: Doc -> Doc
local a = "%" <> a

global :: Doc -> Doc
global a = "@" <> a

label :: Doc -> Doc
label a = "label" <+> "%" <> a

-----
-- Reasoning about types
-----

isFunctionPtr (PointerType (FunctionType {..}) _) = True
isFunctionPtr _ = False

-------------------------------------------------------------------------------
-- Classes
-------------------------------------------------------------------------------

class PP p where
  pp :: p -> Doc

ppMaybe (Just x) = pp x
ppMaybe Nothing = empty

instance PP Word32 where
  pp x = int (fromIntegral x)

instance PP Word64 where
  pp x = int (fromIntegral x)

instance PP Integer where
  pp = integer

instance PP Name where
  pp (Name []) = doubleQuotes empty
  pp (Name name@(first:_))
    | isFirst first && all isRest name = text name
    | otherwise = doubleQuotes . hcat . map escape $ name
    where
        isFirst c = isLetter c || c == '-' || c == '_'
        isRest c = isDigit c || isFirst c
  pp (UnName x) = int (fromIntegral x)

instance PP Parameter where
  pp (Parameter ty name attrs) = pp ty <+> local (pp name)

instance PP ([Parameter], Bool) where
  pp (params, False) = commas (fmap pp params)

instance PP (Operand, [ParameterAttribute]) where
  pp (op, attrs) = ppTyped op

instance PP Type where
  pp (IntegerType width) = "i" <> pp width
  pp (FloatingPointType 16    IEEE) = "half"
  pp (FloatingPointType 32    IEEE) = "float"
  pp (FloatingPointType 64    IEEE) = "double"
  pp (FloatingPointType width IEEE) = "fp" <> pp width
  pp (FloatingPointType width DoubleExtended) = "x86_fp" <> pp width
  pp (FloatingPointType width PairOfFloats)   = "ppc_fp" <> pp width

  pp VoidType = "void"
  pp (PointerType ref@(FunctionType {}) addr) = pp ref
  pp (PointerType ref addr) = pp ref <> "*"
  pp ft@(FunctionType {..}) = pp resultType <+> ppFunctionArgumentTypes ft
  pp (VectorType {..}) = "<" <> pp nVectorElements <+> "x" <+> pp elementType <> ">"
  pp (StructureType {..}) = if isPacked
                               then "<{" <> (commas $ fmap pp elementTypes ) <> "}>"
                               else  "{" <> (commas $ fmap pp elementTypes ) <> "}"
  pp (ArrayType {..}) = brackets $ pp nArrayElements <+> "x" <+> pp elementType
  pp (NamedTypeReference name) = "%" <> pp name

instance PP Global where
  pp (Function {..}) =
      case basicBlocks of
        [] ->
          ("declare" <+> pp returnType <+> global (pp name) <> ppParams (pp . typeOf) parameters)

        -- single unnamed block is special cased, and won't parse otherwise... yeah good times
        [b@(BasicBlock (UnName _) _ _)] ->
            ("define" <+> pp returnType <+> global (pp name) <> ppParams pp parameters)
            `wrapbraces` (nest 2 $ ppSingleBlock b)

        bs ->
          ("define" <+> pp returnType <+> global (pp name) <> ppParams pp parameters)
           `wrapbraces` (vcat $ fmap pp bs)

  pp (GlobalVariable {..}) = global (pp name) <+> "=" <+> "global" <+> pp type' <+> ppMaybe initializer


instance PP Definition where
  pp (GlobalDefinition x) = pp x

instance PP BasicBlock where
  pp (BasicBlock nm instrs term) =
    pp nm <> ":" $+$ nest 2 (vcat $ (fmap pp instrs) ++ [pp term])

instance PP Terminator where
  pp (Br dest meta) = "br" <+> label (pp dest)
  pp (Ret val meta) = "ret" <+> maybe "void" ppTyped val
  pp (CondBr cond tdest fdest meta) =
      "br" <+> ppTyped cond
    <> "," <+> label (pp tdest)
    <> "," <+> label (pp fdest)

instance PP Instruction where
  pp (Mul {..}) = "mul" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (Add {..}) = "add" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (Sub {..}) = "sub" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (FSub {..}) = "fsub" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (FMul {..}) = "fmul" <+> ppTyped operand0 <> "," <+> pp operand1

  pp (FAdd {..}) = "fadd" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (FCmp {..}) = "fcmp" <+> pp fpPredicate <+> ppTyped operand0 <> "," <+> pp operand1

  pp (Alloca {..}) = "alloca" <+> pp allocatedType
  pp (Store {..}) = "store" <+> ppTyped value <> "," <+> ppTyped address
  pp (Load {..}) = "load" <+> ppTyped address
  pp (Phi {..}) = "phi" <+> pp type' <+> commas (fmap phiIncoming incomingValues)

  pp (ICmp {..}) = "icmp" <+> pp iPredicate <+> ppTyped operand0 <> "," <+> pp operand1

  pp c@(Call {..}) = ppCall c

  pp (GetElementPtr {..}) = "getelementptr" <+> bounds inBounds <+> commas (fmap ppTyped (address:indices))
    where
      bounds True = "inbounds"
      bounds False = empty

instance PP CallableOperand where
  pp (Left asm) = undefined
  pp (Right op) = pp op

instance PP Operand where
  pp (LocalReference _ nm) = local (pp nm)
  pp (ConstantOperand con) = pp con


instance PP C.Constant where
  pp (C.Int width val) = pp val
  pp (C.Float (F.Double val)) = text $ printf "%6.6e" val
  pp (C.Float (F.Single val)) = text $ printf "%6.6e" val
  pp (C.GlobalReference ty nm) = "@" <> pp nm

  pp (C.Array {..})
    | memberType == (IntegerType 8) = "c" <> (doubleQuotes $ hcat [ppIntAsChar val | C.Int _ val <- memberValues])
    | otherwise = brackets $ commas $ fmap ppTyped memberValues


  pp (C.GetElementPtr {..}) = "getelementptr" <+> bounds inBounds <+> commas (fmap ppTyped (address:indices))
    where
      bounds True = "inbounds"
      bounds False = empty

instance PP a => PP (Named a) where<commit_msg>support select instruction<commit_after>{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module LLVM.General.Pretty (
  ppllvm,
) where

import GHC.Word

import LLVM.General.Typed

import LLVM.General.AST
import LLVM.General.AST.Global
import LLVM.General.AST.Type

import LLVM.General.AST.Attribute
import qualified LLVM.General.AST.Linkage as L
import qualified LLVM.General.AST.Visibility as V
import qualified LLVM.General.AST.CallingConvention as CC
import qualified LLVM.General.AST.Constant as C
import qualified LLVM.General.AST.FloatingPointPredicate as FP
import qualified LLVM.General.AST.IntegerPredicate as IP
import qualified LLVM.General.AST.AddrSpace as AS
import qualified LLVM.General.AST.Float as F

import Data.String
import Text.Printf
import Text.PrettyPrint

import Data.Char (chr, ord, isControl, isLetter, isDigit)
import Data.List (intersperse)
import Numeric (showHex)

-------------------------------------------------------------------------------
-- Utils
-------------------------------------------------------------------------------

parensIf ::  Bool -> Doc -> Doc
parensIf True = parens
parensIf False = id

commas :: [Doc] -> Doc
commas  = hsep . punctuate (char ',')

colons :: [Doc] -> Doc
colons  = hcat . intersperse (char ':')

hlinecat :: [Doc] -> Doc
hlinecat = vcat . intersperse (text "")

wrapbraces :: Doc -> Doc -> Doc
wrapbraces leadIn x = (leadIn <> char '{') $+$ ( x ) $+$ char '}'

local :: Doc -> Doc
local a = "%" <> a

global :: Doc -> Doc
global a = "@" <> a

label :: Doc -> Doc
label a = "label" <+> "%" <> a

-----
-- Reasoning about types
-----

isFunctionPtr (PointerType (FunctionType {..}) _) = True
isFunctionPtr _ = False

-------------------------------------------------------------------------------
-- Classes
-------------------------------------------------------------------------------

class PP p where
  pp :: p -> Doc

ppMaybe (Just x) = pp x
ppMaybe Nothing = empty

instance PP Word32 where
  pp x = int (fromIntegral x)

instance PP Word64 where
  pp x = int (fromIntegral x)

instance PP Integer where
  pp = integer

instance PP Name where
  pp (Name []) = doubleQuotes empty
  pp (Name name@(first:_))
    | isFirst first && all isRest name = text name
    | otherwise = doubleQuotes . hcat . map escape $ name
    where
        isFirst c = isLetter c || c == '-' || c == '_'
        isRest c = isDigit c || isFirst c
  pp (UnName x) = int (fromIntegral x)

instance PP Parameter where
  pp (Parameter ty name attrs) = pp ty <+> local (pp name)

instance PP ([Parameter], Bool) where
  pp (params, False) = commas (fmap pp params)

instance PP (Operand, [ParameterAttribute]) where
  pp (op, attrs) = ppTyped op

instance PP Type where
  pp (IntegerType width) = "i" <> pp width
  pp (FloatingPointType 16    IEEE) = "half"
  pp (FloatingPointType 32    IEEE) = "float"
  pp (FloatingPointType 64    IEEE) = "double"
  pp (FloatingPointType width IEEE) = "fp" <> pp width
  pp (FloatingPointType width DoubleExtended) = "x86_fp" <> pp width
  pp (FloatingPointType width PairOfFloats)   = "ppc_fp" <> pp width

  pp VoidType = "void"
  pp (PointerType ref@(FunctionType {}) addr) = pp ref
  pp (PointerType ref addr) = pp ref <> "*"
  pp ft@(FunctionType {..}) = pp resultType <+> ppFunctionArgumentTypes ft
  pp (VectorType {..}) = "<" <> pp nVectorElements <+> "x" <+> pp elementType <> ">"
  pp (StructureType {..}) = if isPacked
                               then "<{" <> (commas $ fmap pp elementTypes ) <> "}>"
                               else  "{" <> (commas $ fmap pp elementTypes ) <> "}"
  pp (ArrayType {..}) = brackets $ pp nArrayElements <+> "x" <+> pp elementType
  pp (NamedTypeReference name) = "%" <> pp name

instance PP Global where
  pp (Function {..}) =
      case basicBlocks of
        [] ->
          ("declare" <+> pp returnType <+> global (pp name) <> ppParams (pp . typeOf) parameters)

        -- single unnamed block is special cased, and won't parse otherwise... yeah good times
        [b@(BasicBlock (UnName _) _ _)] ->
            ("define" <+> pp returnType <+> global (pp name) <> ppParams pp parameters)
            `wrapbraces` (nest 2 $ ppSingleBlock b)

        bs ->
          ("define" <+> pp returnType <+> global (pp name) <> ppParams pp parameters)
           `wrapbraces` (vcat $ fmap pp bs)

  pp (GlobalVariable {..}) = global (pp name) <+> "=" <+> "global" <+> pp type' <+> ppMaybe initializer


instance PP Definition where
  pp (GlobalDefinition x) = pp x

instance PP BasicBlock where
  pp (BasicBlock nm instrs term) =
    pp nm <> ":" $+$ nest 2 (vcat $ (fmap pp instrs) ++ [pp term])

instance PP Terminator where
  pp (Br dest meta) = "br" <+> label (pp dest)
  pp (Ret val meta) = "ret" <+> maybe "void" ppTyped val
  pp (CondBr cond tdest fdest meta) =
      "br" <+> ppTyped cond
    <> "," <+> label (pp tdest)
    <> "," <+> label (pp fdest)

instance PP Instruction where
  pp (Mul {..}) = "mul" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (Add {..}) = "add" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (Sub {..}) = "sub" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (FSub {..}) = "fsub" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (FMul {..}) = "fmul" <+> ppTyped operand0 <> "," <+> pp operand1

  pp (FAdd {..}) = "fadd" <+> ppTyped operand0 <> "," <+> pp operand1
  pp (FCmp {..}) = "fcmp" <+> pp fpPredicate <+> ppTyped operand0 <> "," <+> pp operand1

  pp (Alloca {..}) = "alloca" <+> pp allocatedType
  pp (Store {..}) = "store" <+> ppTyped value <> "," <+> ppTyped address
  pp (Load {..}) = "load" <+> ppTyped address
  pp (Phi {..}) = "phi" <+> pp type' <+> commas (fmap phiIncoming incomingValues)

  pp (ICmp {..}) = "icmp" <+> pp iPredicate <+> ppTyped operand0 <> "," <+> pp operand1

  pp c@(Call {..}) = ppCall c
  pp (Select {..}) = "select" <+> pp condition' <+> pp trueValue <+> pp falseValue
  
  pp (GetElementPtr {..}) = "getelementptr" <+> bounds inBounds <+> commas (fmap ppTyped (address:indices))
    where
      bounds True = "inbounds"
      bounds False = empty

instance PP CallableOperand where
  pp (Left asm) = undefined
  pp (Right op) = pp op

instance PP Operand where
  pp (LocalReference _ nm) = local (pp nm)
  pp (ConstantOperand con) = pp con


instance PP C.Constant where
  pp (C.Int width val) = pp val
  pp (C.Float (F.Double val)) = text $ printf "%6.6e" val
  pp (C.Float (F.Single val)) = text $ printf "%6.6e" val
  pp (C.GlobalReference ty nm) = "@" <> pp nm

  pp (C.Array {..})
    | memberType == (IntegerType 8) = "c" <> (doubleQuotes $ hcat [ppIntAsChar val | C.Int _ val <- memberValues])
    | otherwise = brackets $ commas $ fmap ppTyped memberValues


  pp (C.GetElementPtr {..}) = "getelementptr" <+> bounds inBounds <+> commas (fmap ppTyped (address:indices))
    where
      bounds True = "inbounds"
      bounds False = empty

instance PP a => PP (Named a) where
{'fraction_non_alphanumeric': 0.19275466284074605, 'fraction_numerical': 0.005021520803443328, 'mean_word_length': 2.595821511477947, 'pattern_counts': {'":': 2, '<': 178, '<?xml version=': 0, '>': 206, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1652780', 'n_tokens_mistral': 4726, 'n_tokens_neox': 4592, 'n_words': 1462}