-- | Functions to implement the @FR@ (as in "free regs") type class.
--
-- For LLVM GHC calling convention (used registers), see
-- https://github.com/llvm/llvm-project/blob/6ab900f8746e7d8e24afafb5886a40801f6799f4/llvm/lib/Target/RISCV/RISCVISelLowering.cpp#L13638-L13685
module GHC.CmmToAsm.Reg.Linear.RV64
  ( allocateReg,
    getFreeRegs,
    initFreeRegs,
    releaseReg,
    FreeRegs (..),
  )
where

import Data.Word
import GHC.CmmToAsm.RV64.Regs
import GHC.Platform
import GHC.Platform.Reg
import GHC.Platform.Reg.Class.Separate
import GHC.Prelude
import GHC.Stack
import GHC.Utils.Outputable
import GHC.Utils.Panic

-- | Bitmaps to indicate which registers are free (currently unused)
--
-- The bit index represents the `RegNo`, in case of floating point registers
-- with an offset of 32. The register is free when the bit is set.
data FreeRegs
  = FreeRegs
      -- | integer/general purpose registers (`RcInteger`)
      !Word32
      -- | floating point registers (`RcDouble`)
      !Word32

instance Show FreeRegs where
  show :: FreeRegs -> String
show (FreeRegs Word32
g Word32
f) = String
"FreeRegs 0b" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word32 -> String
showBits Word32
g String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" 0b" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word32 -> String
showBits Word32
f

-- | Show bits as a `String` of @1@s and @0@s
showBits :: Word32 -> String
showBits :: Word32 -> String
showBits Word32
w = (Int -> Char) -> [Int] -> String
forall a b. (a -> b) -> [a] -> [b]
map (\Int
i -> if Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
w Int
i then Char
'1' else Char
'0') [Int
0 .. Int
31]

instance Outputable FreeRegs where
  ppr :: FreeRegs -> SDoc
ppr (FreeRegs Word32
g Word32
f) =
    String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"   "
      SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> (Int -> SDoc -> SDoc) -> SDoc -> [Int] -> SDoc
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Int
i SDoc
x -> Int -> SDoc
forall {doc}. IsLine doc => Int -> doc
pad_int Int
i SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> SDoc
x) (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"") [Int
0 .. Int
31]
      SDoc -> SDoc -> SDoc
forall doc. IsDoc doc => doc -> doc -> doc
$$ String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"GPR"
      SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> (Int -> SDoc -> SDoc) -> SDoc -> [Int] -> SDoc
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Int
i SDoc
x -> Word32 -> Int -> SDoc
forall {a} {doc}. (Bits a, IsLine doc) => a -> Int -> doc
show_bit Word32
g Int
i SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> SDoc
x) (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"") [Int
0 .. Int
31]
      SDoc -> SDoc -> SDoc
forall doc. IsDoc doc => doc -> doc -> doc
$$ String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"FPR"
      SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> (Int -> SDoc -> SDoc) -> SDoc -> [Int] -> SDoc
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Int
i SDoc
x -> Word32 -> Int -> SDoc
forall {a} {doc}. (Bits a, IsLine doc) => a -> Int -> doc
show_bit Word32
f Int
i SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> SDoc
x) (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"") [Int
0 .. Int
31]
    where
      pad_int :: Int -> doc
pad_int Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
10 = Char -> doc
forall doc. IsLine doc => Char -> doc
char Char
' ' doc -> doc -> doc
forall doc. IsLine doc => doc -> doc -> doc
<> Int -> doc
forall {doc}. IsLine doc => Int -> doc
int Int
i
      pad_int Int
i = Int -> doc
forall {doc}. IsLine doc => Int -> doc
int Int
i
      -- remember bit = 1 means it's available.
      show_bit :: a -> Int -> doc
show_bit a
bits Int
bit | a -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit a
bits Int
bit = String -> doc
forall doc. IsLine doc => String -> doc
text String
"  "
      show_bit a
_ Int
_ = String -> doc
forall doc. IsLine doc => String -> doc
text String
" x"

-- | Set bits of all allocatable registers to 1
initFreeRegs :: Platform -> FreeRegs
initFreeRegs :: Platform -> FreeRegs
initFreeRegs Platform
platform = (FreeRegs -> RealReg -> FreeRegs)
-> FreeRegs -> [RealReg] -> FreeRegs
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((RealReg -> FreeRegs -> FreeRegs)
-> FreeRegs -> RealReg -> FreeRegs
forall a b c. (a -> b -> c) -> b -> a -> c
flip HasCallStack => RealReg -> FreeRegs -> FreeRegs
RealReg -> FreeRegs -> FreeRegs
releaseReg) FreeRegs
noFreeRegs (Platform -> [RealReg]
allocatableRegs Platform
platform)
  where
    noFreeRegs :: FreeRegs
    noFreeRegs :: FreeRegs
noFreeRegs = Word32 -> Word32 -> FreeRegs
FreeRegs Word32
0 Word32
0

-- | Get all free `RealReg`s (i.e. those where the corresponding bit is 1)
getFreeRegs :: RegClass -> FreeRegs -> [RealReg]
getFreeRegs :: RegClass -> FreeRegs -> [RealReg]
getFreeRegs RegClass
cls (FreeRegs Word32
g Word32
f) =
  case RegClass
cls of
    RegClass
RcInteger -> Int -> Word32 -> [Int] -> [RealReg]
forall {t}. Bits t => Int -> t -> [Int] -> [RealReg]
go Int
0 Word32
g [Int]
allocatableIntRegs
    RegClass
RcFloat -> Int -> Word32 -> [Int] -> [RealReg]
forall {t}. Bits t => Int -> t -> [Int] -> [RealReg]
go Int
32 Word32
f [Int]
allocatableDoubleRegs
    RegClass
RcVector ->
      String -> [RealReg]
forall a. HasCallStack => String -> a
sorry String
"Linear.RV64.getFreeRegs: vector registers are not supported"

  where
    go :: Int -> t -> [Int] -> [RealReg]
go Int
_ t
_ [] = []
    go Int
off t
x (Int
i : [Int]
is)
      | t -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit t
x Int
i = Int -> RealReg
RealRegSingle (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) RealReg -> [RealReg] -> [RealReg]
forall a. a -> [a] -> [a]
: (Int -> t -> [Int] -> [RealReg]
go Int
off t
x ([Int] -> [RealReg]) -> [Int] -> [RealReg]
forall a b. (a -> b) -> a -> b
$! [Int]
is)
      | Bool
otherwise = Int -> t -> [Int] -> [RealReg]
go Int
off t
x ([Int] -> [RealReg]) -> [Int] -> [RealReg]
forall a b. (a -> b) -> a -> b
$! [Int]
is
    -- The lists of allocatable registers are manually crafted: Register
    -- allocation is pretty hot code. We don't want to iterate and map like
    -- `initFreeRegs` all the time! (The register mappings aren't supposed to
    -- change often.)
    allocatableIntRegs :: [Int]
allocatableIntRegs = [Int
5 .. Int
7] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
10 .. Int
17] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
28 .. Int
30]
    allocatableDoubleRegs :: [Int]
allocatableDoubleRegs = [Int
0 .. Int
7] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
10 .. Int
17] [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int
28 .. Int
31]

-- | Set corresponding register bit to 0
allocateReg :: (HasCallStack) => RealReg -> FreeRegs -> FreeRegs
allocateReg :: HasCallStack => RealReg -> FreeRegs -> FreeRegs
allocateReg (RealRegSingle Int
r) (FreeRegs Word32
g Word32
f)
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
31 Bool -> Bool -> Bool
&& Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
f (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
32) = Word32 -> Word32 -> FreeRegs
FreeRegs Word32
g (Word32 -> Int -> Word32
forall a. (Num a, Bits a) => a -> Int -> a
clearBit Word32
f (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
32))
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
32 Bool -> Bool -> Bool
&& Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
g Int
r = Word32 -> Word32 -> FreeRegs
FreeRegs (Word32 -> Int -> Word32
forall a. (Num a, Bits a) => a -> Int -> a
clearBit Word32
g Int
r) Word32
f
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
31 = String -> FreeRegs
forall a. HasCallStack => String -> a
panic (String -> FreeRegs) -> String -> FreeRegs
forall a b. (a -> b) -> a -> b
$ String
"Linear.RV64.allocReg: double allocation of float reg v" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
32) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"; " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word32 -> String
showBits Word32
f
  | Bool
otherwise = String -> SDoc -> FreeRegs
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"Linear.RV64.allocReg" (SDoc -> FreeRegs) -> SDoc -> FreeRegs
forall a b. (a -> b) -> a -> b
$ String -> SDoc
forall doc. IsLine doc => String -> doc
text (String
"double allocation of gp reg x" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
r String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"; " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Word32 -> String
showBits Word32
g)

-- | Set corresponding register bit to 1
releaseReg :: (HasCallStack) => RealReg -> FreeRegs -> FreeRegs
releaseReg :: HasCallStack => RealReg -> FreeRegs -> FreeRegs
releaseReg (RealRegSingle Int
r) (FreeRegs Word32
g Word32
f)
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
31 Bool -> Bool -> Bool
&& Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
f (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
32) = String -> SDoc -> FreeRegs
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"Linear.RV64.releaseReg" (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"can't release non-allocated reg v" SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Int -> SDoc
forall {doc}. IsLine doc => Int -> doc
int (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
32))
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
32 Bool -> Bool -> Bool
&& Word32 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word32
g Int
r = String -> SDoc -> FreeRegs
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"Linear.RV64.releaseReg" (String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"can't release non-allocated reg x" SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Int -> SDoc
forall {doc}. IsLine doc => Int -> doc
int Int
r)
  | Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
31 = Word32 -> Word32 -> FreeRegs
FreeRegs Word32
g (Word32 -> Int -> Word32
forall a. (Num a, Bits a) => a -> Int -> a
setBit Word32
f (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
32))
  | Bool
otherwise = Word32 -> Word32 -> FreeRegs
FreeRegs (Word32 -> Int -> Word32
forall a. (Num a, Bits a) => a -> Int -> a
setBit Word32
g Int
r) Word32
f