-- | The frame table, one per JIT. A native caller whose callee exits writes a
-- frame record naming an entry here; the trampoline turns the record into
-- the @Push@ frame the interpreter would have pushed. See unison-runtime/src/Unison/Runtime/JIT/design.md,
-- "Frames: native code records, Haskell builds".
module Unison.Runtime.JIT.Frames
  ( Frame (..),
    FrameIndex,
    Frames,
    newFrames,
    registerFrames,
    replaceFrames,
    lookupFrame,
  )
where

import Control.Monad (forM_)
import Control.Monad.Primitive (PrimState)
import Data.IORef
import Data.Primitive.SmallArray (SmallMutableArray, copySmallMutableArray, getSizeofSmallMutableArray, newSmallArray, readSmallArray, writeSmallArray)
import Foreign.Ptr (Ptr)
import Unison.Runtime.MCode (CombIx, NativeCell)
import Unison.Runtime.Stack (MSection)

-- | The fixed parts of a @Push@ frame for one @Let@: the body's @CombIx@,
-- stack guard and section, and the cell holding the body's native code.
data Frame = Frame !CombIx !Int !MSection !(Ptr NativeCell)

-- | 1-based.
type FrameIndex = Int

-- | Like the exits table: an array that grows by copying, written only by
-- the thread that compiles.
data Table = Table !Int !(SmallMutableArray (PrimState IO) Frame)

newtype Frames = Frames (IORef Table)

newFrames :: IO Frames
newFrames :: IO Frames
newFrames = IORef Table -> Frames
Frames (IORef Table -> Frames) -> IO (IORef Table) -> IO Frames
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Table -> IO (IORef Table)
forall a. a -> IO (IORef a)
newIORef (Table -> IO (IORef Table))
-> (SmallMutableArray RealWorld Frame -> Table)
-> SmallMutableArray RealWorld Frame
-> IO (IORef Table)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FrameIndex -> SmallMutableArray (PrimState IO) Frame -> Table
Table FrameIndex
1 (SmallMutableArray RealWorld Frame -> IO (IORef Table))
-> IO (SmallMutableArray RealWorld Frame) -> IO (IORef Table)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< FrameIndex -> Frame -> IO (SmallMutableArray (PrimState IO) Frame)
forall (m :: * -> *) a.
PrimMonad m =>
FrameIndex -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray FrameIndex
1024 Frame
unknown)

unknown :: Frame
unknown :: Frame
unknown = [Char] -> Frame
forall a. HasCallStack => [Char] -> a
error [Char]
"JIT: unknown frame index"

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

-- | Overwrites frames registered earlier, starting at the given index
-- (see replaceExits).
replaceFrames :: Frames -> FrameIndex -> [Frame] -> IO ()
replaceFrames :: Frames -> FrameIndex -> [Frame] -> IO ()
replaceFrames (Frames IORef Table
table) FrameIndex
base [Frame]
frames = do
  Table FrameIndex
_ SmallMutableArray (PrimState IO) Frame
arr <- IORef Table -> IO Table
forall a. IORef a -> IO a
readIORef IORef Table
table
  [(FrameIndex, Frame)] -> ((FrameIndex, Frame) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ ([FrameIndex] -> [Frame] -> [(FrameIndex, Frame)]
forall a b. [a] -> [b] -> [(a, b)]
zip [FrameIndex
base ..] [Frame]
frames) (((FrameIndex, Frame) -> IO ()) -> IO ())
-> ((FrameIndex, Frame) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(FrameIndex
i, Frame
f) -> SmallMutableArray (PrimState IO) Frame
-> FrameIndex -> Frame -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> FrameIndex -> a -> m ()
writeSmallArray SmallMutableArray (PrimState IO) Frame
arr FrameIndex
i (Frame -> IO ()) -> Frame -> IO ()
forall a b. (a -> b) -> a -> b
$! Frame
f

lookupFrame :: Frames -> FrameIndex -> IO Frame
lookupFrame :: Frames -> FrameIndex -> IO Frame
lookupFrame (Frames IORef Table
table) FrameIndex
i = do
  Table FrameIndex
_ SmallMutableArray (PrimState IO) Frame
arr <- IORef Table -> IO Table
forall a. IORef a -> IO a
readIORef IORef Table
table
  SmallMutableArray (PrimState IO) Frame -> FrameIndex -> IO Frame
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> FrameIndex -> m a
readSmallArray SmallMutableArray (PrimState IO) Frame
arr FrameIndex
i