{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

-- | The exits table, one per JIT (which is one per process, see JIT.hs). Native code returns an index into it to say
-- what the interpreter should do next. See unison-runtime/src/Unison/Runtime/JIT/design.md.
module Unison.Runtime.JIT.Exits
  ( Exit (..),
    isAdHoc,
    ExitIndex,
    Exits,
    newExits,
    registerExits,
    replaceExits,
    lookupExit,
    countExit,
    exitCounts,
    resetExitCounts,
    countEntry,
    nameEntry,
    entryCounts,
  )
where

import Control.Monad (forM, forM_)
import Control.Monad.Primitive (PrimState)
import Data.IORef
import Data.Map.Strict qualified as M
import Data.Primitive.PrimArray (MutablePrimArray (..), newPrimArray, readPrimArray, setPrimArray)
import Data.Primitive.SmallArray (SmallMutableArray, copySmallMutableArray, getSizeofSmallMutableArray, newSmallArray, readSmallArray, writeSmallArray)
import Foreign.Ptr (Ptr, WordPtr, ptrToWordPtr)
import GHC.Exts (Int (..), fetchAddIntArray#)
import GHC.IO (IO (..))
import Unison.Runtime.MCode (CombIx, NativeCell)
import Unison.Runtime.Stack (MInstr, MSection)

data Exit
  = -- | interpret from this section of this combinator
    Resume !CombIx !MSection
  | -- | grow the Unison stack to fit a frame of this size, then call the
    -- function whose cell this is again (it may not be the one the
    -- trampoline entered: tail calls change functions)
    GrowStack !Int !(Ptr NativeCell)
  | -- | the runtime asked the thread to stop; call the function whose
    -- cell this is again when the thread next runs
    Reenter !(Ptr NativeCell)
  | -- | run this instruction with the interpreter (it pushes the given
    -- number of values), then call the function in this cell, which
    -- continues with the section after it. If the instruction pushed a
    -- different number of values, interpret that section instead.
    CallOut !CombIx !MInstr !MSection !Int !(Ptr NativeCell)
  | -- | a description, for statistics
    Named String Exit

-- | An exit the code didn't cause: the poll (the context-switch flag or
-- the allocation budget) and stack growth land wherever the program
-- happens to be. The others, a call-out or a resume, are taken at a point
-- the code itself reaches on every such path.
isAdHoc :: Exit -> Bool
isAdHoc :: Exit -> Bool
isAdHoc (GrowStack Int
_ Ptr NativeCell
_) = Bool
True
isAdHoc (Reenter Ptr NativeCell
_) = Bool
True
isAdHoc (Named String
_ Exit
e) = Exit -> Bool
isAdHoc Exit
e
isAdHoc Exit
_ = Bool
False

-- | 1-based; 0 is OK and negative values are errors.
type ExitIndex = Int

-- | The exits, by index, in an array that grows by copying. Looking one up
-- is on the path of every exit, so it is two loads. Only the thread that
-- compiles writes; an entry is in place before the code that returns its
-- index is installed, and a reader that still holds an older array finds
-- every index it can be handed there.
data Table = Table !Int !(SmallMutableArray (PrimState IO) Exit) -- next index, exits

-- | The table, the exit counters, and the entry counts (by code address)
-- with the names of the functions entered.
data Exits = Exits
  { Exits -> IORef Table
exTable :: !(IORef Table),
    Exits -> MutablePrimArray (PrimState IO) Int
exCounters :: !(MutablePrimArray (PrimState IO) Int),
    Exits -> IORef (Map WordPtr Int)
exEntries :: !(IORef (M.Map WordPtr Int)),
    Exits -> IORef (Map WordPtr String)
exEntryNames :: !(IORef (M.Map WordPtr String))
  }

-- | An empty table. With statistics on, the counters have a slot per exit
-- index (up to 'counterSlots'); otherwise one, since nothing counts.
newExits :: Bool -> IO Exits
newExits :: Bool -> IO Exits
newExits Bool
stats = do
  IORef Table
t <- Table -> IO (IORef Table)
forall a. a -> IO (IORef a)
newIORef (Table -> IO (IORef Table))
-> (SmallMutableArray RealWorld Exit -> Table)
-> SmallMutableArray RealWorld Exit
-> IO (IORef Table)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> SmallMutableArray (PrimState IO) Exit -> Table
Table Int
1 (SmallMutableArray RealWorld Exit -> IO (IORef Table))
-> IO (SmallMutableArray RealWorld Exit) -> IO (IORef Table)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Int -> Exit -> IO (SmallMutableArray (PrimState IO) Exit)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
1024 Exit
unknown
  let n :: Int
n = if Bool
stats then Int
counterSlots else Int
1
  MutablePrimArray RealWorld Int
arr <- Int -> IO (MutablePrimArray (PrimState IO) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
n
  MutablePrimArray (PrimState IO) Int -> Int -> Int -> Int -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
setPrimArray MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
arr Int
0 Int
n Int
0
  IORef Table
-> MutablePrimArray (PrimState IO) Int
-> IORef (Map WordPtr Int)
-> IORef (Map WordPtr String)
-> Exits
Exits IORef Table
t MutablePrimArray RealWorld Int
MutablePrimArray (PrimState IO) Int
arr (IORef (Map WordPtr Int) -> IORef (Map WordPtr String) -> Exits)
-> IO (IORef (Map WordPtr Int))
-> IO (IORef (Map WordPtr String) -> Exits)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map WordPtr Int -> IO (IORef (Map WordPtr Int))
forall a. a -> IO (IORef a)
newIORef Map WordPtr Int
forall k a. Map k a
M.empty IO (IORef (Map WordPtr String) -> Exits)
-> IO (IORef (Map WordPtr String)) -> IO Exits
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map WordPtr String -> IO (IORef (Map WordPtr String))
forall a. a -> IO (IORef a)
newIORef Map WordPtr String
forall k a. Map k a
M.empty

unknown :: Exit
unknown :: Exit
unknown = String -> Exit
forall a. HasCallStack => String -> a
error String
"JIT: unknown exit index"

-- | Adds a module's exits and returns the index of the first one. The
-- rest follow consecutively, in the order given.
registerExits :: Exits -> [Exit] -> IO ExitIndex
registerExits :: Exits -> [Exit] -> IO Int
registerExits Exits
ex [Exit]
exits = do
  Table Int
next SmallMutableArray (PrimState IO) Exit
arr <- IORef Table -> IO Table
forall a. IORef a -> IO a
readIORef (Exits -> IORef Table
exTable Exits
ex)
  Int
size <- SmallMutableArray (PrimState IO) Exit -> IO Int
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m Int
getSizeofSmallMutableArray SmallMutableArray (PrimState IO) Exit
arr
  let n :: Int
n = [Exit] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Exit]
exits
  SmallMutableArray RealWorld Exit
arr' <-
    if Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
size
      then SmallMutableArray RealWorld Exit
-> IO (SmallMutableArray RealWorld Exit)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray RealWorld Exit
SmallMutableArray (PrimState IO) Exit
arr
      else do
        SmallMutableArray RealWorld Exit
bigger <- Int -> Exit -> IO (SmallMutableArray (PrimState IO) Exit)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
size) (Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)) Exit
unknown
        SmallMutableArray (PrimState IO) Exit
-> Int
-> SmallMutableArray (PrimState IO) Exit
-> Int
-> Int
-> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray RealWorld Exit
SmallMutableArray (PrimState IO) Exit
bigger Int
0 SmallMutableArray (PrimState IO) Exit
arr Int
0 Int
next
        pure SmallMutableArray RealWorld Exit
bigger
  [(Int, Exit)] -> ((Int, Exit) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [Exit] -> [(Int, Exit)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
next ..] [Exit]
exits) (((Int, Exit) -> IO ()) -> IO ())
-> ((Int, Exit) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, Exit
e) -> SmallMutableArray (PrimState IO) Exit -> Int -> Exit -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld Exit
SmallMutableArray (PrimState IO) Exit
arr' Int
i (Exit -> IO ()) -> Exit -> IO ()
forall a b. (a -> b) -> a -> b
$! Exit
e
  IORef Table -> Table -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Exits -> IORef Table
exTable Exits
ex) (Int -> SmallMutableArray (PrimState IO) Exit -> Table
Table (Int
next Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) SmallMutableArray RealWorld Exit
SmallMutableArray (PrimState IO) Exit
arr')
  pure Int
next

-- | Overwrites exits registered earlier, starting at the given index: the
-- second code generation pass produces the final ones (with the cells of
-- the auxiliary functions), the first only their number.
replaceExits :: Exits -> ExitIndex -> [Exit] -> IO ()
replaceExits :: Exits -> Int -> [Exit] -> IO ()
replaceExits Exits
ex Int
base [Exit]
exits = do
  Table Int
_ SmallMutableArray (PrimState IO) Exit
arr <- IORef Table -> IO Table
forall a. IORef a -> IO a
readIORef (Exits -> IORef Table
exTable Exits
ex)
  [(Int, Exit)] -> ((Int, Exit) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([Int] -> [Exit] -> [(Int, Exit)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
base ..] [Exit]
exits) (((Int, Exit) -> IO ()) -> IO ())
-> ((Int, Exit) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Int
i, Exit
e) -> SmallMutableArray (PrimState IO) Exit -> Int -> Exit -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState IO) Exit
arr Int
i (Exit -> IO ()) -> Exit -> IO ()
forall a b. (a -> b) -> a -> b
$! Exit
e

lookupExit :: Exits -> ExitIndex -> IO Exit
lookupExit :: Exits -> Int -> IO Exit
lookupExit Exits
ex Int
i = do
  Table Int
_ SmallMutableArray (PrimState IO) Exit
arr <- IORef Table -> IO Table
forall a. IORef a -> IO a
readIORef (Exits -> IORef Table
exTable Exits
ex)
  SmallMutableArray (PrimState IO) Exit -> Int -> IO Exit
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray (PrimState IO) Exit
arr Int
i
{-# INLINE lookupExit #-}

-- | How often each exit was taken, for @UNISON_JIT_STATS@: one counter per
-- exit index, with the total in slot 0. Fixed in size so that counting is
-- one atomic add; exits past the end are counted together in the last
-- slot. Only allocated when statistics are on.
counterSlots :: Int
counterSlots :: Int
counterSlots = Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024

-- | Counts an exit; gives the total taken so far.
countExit :: Exits -> ExitIndex -> IO Int
countExit :: Exits -> Int -> IO Int
countExit Exits
ex Int
i = do
  Int
_ <- MutablePrimArray (PrimState IO) Int -> Int -> Int -> IO Int
fetchAddInt (Exits -> MutablePrimArray (PrimState IO) Int
exCounters Exits
ex) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
i (Int
counterSlots Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Int
1
  (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Int) -> IO Int -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState IO) Int -> Int -> Int -> IO Int
fetchAddInt (Exits -> MutablePrimArray (PrimState IO) Int
exCounters Exits
ex) Int
0 Int
1

-- | Every exit that has been taken, with its count.
exitCounts :: Exits -> IO [(ExitIndex, Exit, Int)]
exitCounts :: Exits -> IO [(Int, Exit, Int)]
exitCounts Exits
ex = do
  Table Int
next SmallMutableArray (PrimState IO) Exit
arr <- IORef Table -> IO Table
forall a. IORef a -> IO a
readIORef (Exits -> IORef Table
exTable Exits
ex)
  [(Int, Int)]
counts <- [Int] -> (Int -> IO (Int, Int)) -> IO [(Int, Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Int
1 .. Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
next Int
counterSlots Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ((Int -> IO (Int, Int)) -> IO [(Int, Int)])
-> (Int -> IO (Int, Int)) -> IO [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ \Int
i -> (,) Int
i (Int -> (Int, Int)) -> IO Int -> IO (Int, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState IO) Int -> Int -> IO Int
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
readPrimArray (Exits -> MutablePrimArray (PrimState IO) Int
exCounters Exits
ex) Int
i
  [(Int, Int)]
-> ((Int, Int) -> IO (Int, Exit, Int)) -> IO [(Int, Exit, Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Int
i, Int
n) | (Int
i, Int
n) <- [(Int, Int)]
counts, Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0] (((Int, Int) -> IO (Int, Exit, Int)) -> IO [(Int, Exit, Int)])
-> ((Int, Int) -> IO (Int, Exit, Int)) -> IO [(Int, Exit, Int)]
forall a b. (a -> b) -> a -> b
$ \(Int
i, Int
n) -> (\Exit
e -> (Int
i, Exit
e, Int
n)) (Exit -> (Int, Exit, Int)) -> IO Exit -> IO (Int, Exit, Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallMutableArray (PrimState IO) Exit -> Int -> IO Exit
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray (PrimState IO) Exit
arr Int
i

-- | Starts counting again from zero.
resetExitCounts :: Exits -> IO ()
resetExitCounts :: Exits -> IO ()
resetExitCounts Exits
ex = do
  Table Int
next SmallMutableArray (PrimState IO) Exit
_ <- IORef Table -> IO Table
forall a. IORef a -> IO a
readIORef (Exits -> IORef Table
exTable Exits
ex)
  MutablePrimArray (PrimState IO) Int -> Int -> Int -> Int -> IO ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
setPrimArray (Exits -> MutablePrimArray (PrimState IO) Int
exCounters Exits
ex) Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
next Int
counterSlots) Int
0

fetchAddInt :: MutablePrimArray (PrimState IO) Int -> Int -> Int -> IO Int
fetchAddInt :: MutablePrimArray (PrimState IO) Int -> Int -> Int -> IO Int
fetchAddInt (MutablePrimArray MutableByteArray# (PrimState IO)
arr) (I# Int#
i) (I# Int#
n) = (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO ((State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int)
-> (State# RealWorld -> (# State# RealWorld, Int #)) -> IO Int
forall a b. (a -> b) -> a -> b
$ \State# RealWorld
s -> case MutableByteArray# RealWorld
-> Int# -> Int# -> State# RealWorld -> (# State# RealWorld, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# RealWorld
MutableByteArray# (PrimState IO)
arr Int#
i Int#
n State# RealWorld
s of
  (# State# RealWorld
s', Int#
old #) -> (# State# RealWorld
s', Int# -> Int
I# Int#
old #)

-- | For @UNISON_JIT_STATS@: how often the trampoline entered each native
-- function (by code address), and what each is called. Slow (a map in an
-- IORef); statistics runs only.
countEntry :: Exits -> Ptr a -> IO ()
countEntry :: forall a. Exits -> Ptr a -> IO ()
countEntry Exits
ex Ptr a
fn = IORef (Map WordPtr Int)
-> (Map WordPtr Int -> (Map WordPtr Int, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Exits -> IORef (Map WordPtr Int)
exEntries Exits
ex) (\Map WordPtr Int
m -> ((Int -> Int -> Int)
-> WordPtr -> Int -> Map WordPtr Int -> Map WordPtr Int
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) (Ptr a -> WordPtr
forall a. Ptr a -> WordPtr
ptrToWordPtr Ptr a
fn) Int
1 Map WordPtr Int
m, ()))

-- | Records the name of the function at a code address.
nameEntry :: Exits -> Ptr a -> String -> IO ()
nameEntry :: forall a. Exits -> Ptr a -> String -> IO ()
nameEntry Exits
ex Ptr a
fn String
name = IORef (Map WordPtr String)
-> (Map WordPtr String -> (Map WordPtr String, ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Exits -> IORef (Map WordPtr String)
exEntryNames Exits
ex) (\Map WordPtr String
m -> (WordPtr -> String -> Map WordPtr String -> Map WordPtr String
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert (Ptr a -> WordPtr
forall a. Ptr a -> WordPtr
ptrToWordPtr Ptr a
fn) String
name Map WordPtr String
m, ()))

-- | The entries counted since the last call, by function name.
entryCounts :: Exits -> IO [(String, Int)]
entryCounts :: Exits -> IO [(String, Int)]
entryCounts Exits
ex = do
  Map WordPtr Int
counts <- IORef (Map WordPtr Int)
-> (Map WordPtr Int -> (Map WordPtr Int, Map WordPtr Int))
-> IO (Map WordPtr Int)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Exits -> IORef (Map WordPtr Int)
exEntries Exits
ex) (\Map WordPtr Int
m -> (Map WordPtr Int
forall k a. Map k a
M.empty, Map WordPtr Int
m))
  Map WordPtr String
names <- IORef (Map WordPtr String) -> IO (Map WordPtr String)
forall a. IORef a -> IO a
readIORef (Exits -> IORef (Map WordPtr String)
exEntryNames Exits
ex)
  pure [(String -> WordPtr -> Map WordPtr String -> String
forall k a. Ord k => a -> k -> Map k a -> a
M.findWithDefault (WordPtr -> String
forall a. Show a => a -> String
show WordPtr
a) WordPtr
a Map WordPtr String
names, Int
n) | (WordPtr
a, Int
n) <- Map WordPtr Int -> [(WordPtr, Int)]
forall k a. Map k a -> [(k, a)]
M.toList Map WordPtr Int
counts]