-- | The constant pool: Haskell heap objects that native code needs, in an
-- array whose address is passed on every entry. See unison-runtime/src/Unison/Runtime/JIT/design.md,
-- "Appendix: the constant pool", which also says why there is one pool per
-- process rather than one per module.
--
-- Indices are assigned when a module is compiled ('poolIndices') and are
-- constants in the generated code. The array grows by copying; the old
-- arrays are kept, and new entries are written to them too where they
-- fit, so an address read at entry stays valid for as long as that native
-- run lasts. A run can still meet code that was installed after it began
-- and uses an index past the end of the array the run holds: a function
-- that uses an index of 'poolStableSize' or more checks the array's size
-- at entry, and exits to be entered again with the current array.
module Unison.Runtime.JIT.Pool
  ( Pool,
    newPool,
    PoolKey (..),
    poolIndices,
    currentPool,
    poolStableSize,
    poolIndexCharTag,
    poolIndexFloatTag,
    poolIndexIntTag,
    poolIndexNatTag,
    poolIndexTrue,
    poolIndexFalse,
  )
where

import Control.Monad (forM, forM_, when)
import Control.Monad.Primitive (PrimState)
import Data.IORef
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import Data.Primitive.Array (MutableArray, copyMutableArray, newArray, sizeofMutableArray, writeArray)
import Data.Vector qualified as V
import Unison.Reference (Reference)
import Unison.Runtime.ANF (PackedTag)
import Unison.Runtime.MCode (CombIx, GCombInfo, MLit (..), RecordShape (..))
import Unison.Runtime.Stack (Closure, MComb, Val (..), charTypeTag, falseVal, floatTypeTag, intTypeTag, natTypeTag, nullSeg, trueVal, pattern Enum, pattern Foreign, pattern PAp, pattern RecordG)
import Unison.Runtime.Stack qualified as Stack

-- The order matches unboxedTypeTagToInt in Stack.hs.
poolIndexCharTag, poolIndexFloatTag, poolIndexIntTag, poolIndexNatTag :: Int
poolIndexCharTag :: Int
poolIndexCharTag = Int
0
poolIndexFloatTag :: Int
poolIndexFloatTag = Int
1
poolIndexIntTag :: Int
poolIndexIntTag = Int
2
poolIndexNatTag :: Int
poolIndexNatTag = Int
3

-- | The boolean closures, for comparison results.
poolIndexTrue, poolIndexFalse :: Int
poolIndexTrue :: Int
poolIndexTrue = Int
4
poolIndexFalse :: Int
poolIndexFalse = Int
5

fixedEntries :: Int
fixedEntries :: Int
fixedEntries = Int
6

-- | The size of the first array: every array ever passed to native code
-- has at least this many entries.
-- | A pool, with the size its first array has: a function that uses only
-- indices below it never has to check the array it was given (the stress
-- mode @pool=N@ makes it tiny, so that growing is exercised).
data Pool = Pool !(IORef PoolState) !Int

poolStableSize :: Pool -> Int
poolStableSize :: Pool -> Int
poolStableSize (Pool IORef PoolState
_ Int
n) = Int
n

-- | What a pool entry is. Entries are interned by key.
data PoolKey
  = -- | a constructor with no fields. Also how a 'Reference' is kept: the
    -- entry for @KeyEnum r 0@ is an 'Enum' whose reference field native
    -- code reads when it builds a constructor of that type.
    KeyEnum !Reference !PackedTag
  | -- | a boxed literal (text, term link, type link)
    KeyLit !MLit
  | -- | a known combinator as a value: a @PAp@ with nothing captured, the
    -- same closure the interpreter builds for @App (Env cix) ZArgs@.
    -- Interned by the CombIx alone (the combinator has no Ord).
    KeyComb !CombIx !(GCombInfo MComb)
  | -- | the boxed part of a top-level value that was evaluated when its
    -- definition was loaded (a @CachedVal@), by its CombIx
    KeyCached !CombIx !Closure
  | -- | a record shape, kept as an empty record of that shape: native code
    -- copies the shape's words out of it when it builds a record
    -- (@RecPack@). Interned by the field names, which determine the shape.
    KeyRecord !RecordShape

instance Eq PoolKey where
  PoolKey
a == :: PoolKey -> PoolKey -> Bool
== PoolKey
b = PoolKey -> PoolKey -> Ordering
forall a. Ord a => a -> a -> Ordering
compare PoolKey
a PoolKey
b Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ

instance Ord PoolKey where
  compare :: PoolKey -> PoolKey -> Ordering
compare = \cases
    (KeyEnum Reference
r PackedTag
t) (KeyEnum Reference
r' PackedTag
t') -> (Reference, PackedTag) -> (Reference, PackedTag) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Reference
r, PackedTag
t) (Reference
r', PackedTag
t')
    (KeyEnum {}) PoolKey
_ -> Ordering
LT
    PoolKey
_ (KeyEnum {}) -> Ordering
GT
    (KeyLit MLit
l) (KeyLit MLit
l') -> MLit -> MLit -> Ordering
forall a. Ord a => a -> a -> Ordering
compare MLit
l MLit
l'
    (KeyLit {}) PoolKey
_ -> Ordering
LT
    PoolKey
_ (KeyLit {}) -> Ordering
GT
    (KeyComb CombIx
c GCombInfo MComb
_) (KeyComb CombIx
c' GCombInfo MComb
_) -> CombIx -> CombIx -> Ordering
forall a. Ord a => a -> a -> Ordering
compare CombIx
c CombIx
c'
    (KeyComb {}) PoolKey
_ -> Ordering
LT
    PoolKey
_ (KeyComb {}) -> Ordering
GT
    (KeyCached CombIx
c Closure
_) (KeyCached CombIx
c' Closure
_) -> CombIx -> CombIx -> Ordering
forall a. Ord a => a -> a -> Ordering
compare CombIx
c CombIx
c'
    (KeyCached {}) PoolKey
_ -> Ordering
LT
    PoolKey
_ (KeyCached {}) -> Ordering
GT
    (KeyRecord RecordShape
s) (KeyRecord RecordShape
s') -> RecordShape -> RecordShape -> Ordering
forall a. Ord a => a -> a -> Ordering
compare RecordShape
s RecordShape
s'

instance Show PoolKey where
  show :: PoolKey -> String
show = \case
    KeyEnum Reference
r PackedTag
t -> String
"KeyEnum " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Reference -> String
forall a. Show a => a -> String
show Reference
r String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ PackedTag -> String
forall a. Show a => a -> String
show PackedTag
t
    KeyLit MLit
l -> String
"KeyLit " String -> ShowS
forall a. [a] -> [a] -> [a]
++ MLit -> String
forall a. Show a => a -> String
show MLit
l
    KeyComb CombIx
c GCombInfo MComb
_ -> String
"KeyComb " String -> ShowS
forall a. [a] -> [a] -> [a]
++ CombIx -> String
forall a. Show a => a -> String
show CombIx
c
    KeyCached CombIx
c Closure
_ -> String
"KeyCached " String -> ShowS
forall a. [a] -> [a] -> [a]
++ CombIx -> String
forall a. Show a => a -> String
show CombIx
c
    KeyRecord RecordShape
s -> String
"KeyRecord " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Vector Text -> String
forall a. Show a => a -> String
show (RecordShape -> Vector Text
shapeFields RecordShape
s)

data PoolState = PoolState
  { PoolState -> MutableArray (PrimState IO) Closure
poolArray :: !(MutableArray (PrimState IO) Closure),
    PoolState -> Int
poolNext :: !Int,
    PoolState -> Map PoolKey Int
poolKeys :: !(Map.Map PoolKey Int),
    -- | superseded arrays, kept alive on purpose
    PoolState -> [MutableArray (PrimState IO) Closure]
poolOld :: [MutableArray (PrimState IO) Closure]
  }

-- | A pool holding the fixed entries. The argument is the stress mode's
-- size for the first array, if any.
newPool :: Maybe Int -> IO Pool
newPool :: Maybe Int -> IO Pool
newPool Maybe Int
stress = do
  let stable :: Int
stable = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
fixedEntries (Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
4096 Maybe Int
stress)
  -- Everything is forced before it goes in: native code reads these
  -- objects' fields directly, so they must be values, not thunks.
  MutableArray RealWorld Closure
arr <- Int -> Closure -> IO (MutableArray (PrimState IO) Closure)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
stable (Closure -> IO (MutableArray RealWorld Closure))
-> Closure -> IO (MutableArray RealWorld Closure)
forall a b. (a -> b) -> a -> b
$! Closure
natTypeTag
  let put :: Int -> Closure -> IO ()
put Int
i Closure
c = MutableArray (PrimState IO) Closure -> Int -> Closure -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld Closure
MutableArray (PrimState IO) Closure
arr Int
i (Closure -> IO ()) -> Closure -> IO ()
forall a b. (a -> b) -> a -> b
$! Closure
c
  Int -> Closure -> IO ()
put Int
poolIndexCharTag Closure
charTypeTag
  Int -> Closure -> IO ()
put Int
poolIndexFloatTag Closure
floatTypeTag
  Int -> Closure -> IO ()
put Int
poolIndexIntTag Closure
intTypeTag
  Int -> Closure -> IO ()
put Int
poolIndexNatTag Closure
natTypeTag
  Int -> Closure -> IO ()
put Int
poolIndexTrue (Val -> Closure
getBoxedVal Val
trueVal)
  Int -> Closure -> IO ()
put Int
poolIndexFalse (Val -> Closure
getBoxedVal Val
falseVal)
  IORef PoolState
ref <- PoolState -> IO (IORef PoolState)
forall a. a -> IO (IORef a)
newIORef (MutableArray (PrimState IO) Closure
-> Int
-> Map PoolKey Int
-> [MutableArray (PrimState IO) Closure]
-> PoolState
PoolState MutableArray RealWorld Closure
MutableArray (PrimState IO) Closure
arr Int
fixedEntries Map PoolKey Int
forall k a. Map k a
Map.empty [])
  pure (IORef PoolState -> Int -> Pool
Pool IORef PoolState
ref Int
stable)

-- | The array to pass to native code now.
currentPool :: Pool -> IO (MutableArray (PrimState IO) Closure)
currentPool :: Pool -> IO (MutableArray (PrimState IO) Closure)
currentPool (Pool IORef PoolState
ref Int
_) = PoolState -> MutableArray RealWorld Closure
PoolState -> MutableArray (PrimState IO) Closure
poolArray (PoolState -> MutableArray RealWorld Closure)
-> IO PoolState -> IO (MutableArray RealWorld Closure)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef PoolState -> IO PoolState
forall a. IORef a -> IO a
readIORef IORef PoolState
ref

-- | The index of each key, adding the ones not seen before. Called when
-- a module is compiled, before its code is generated.
poolIndices :: Pool -> [PoolKey] -> IO (Map.Map PoolKey Int)
poolIndices :: Pool -> [PoolKey] -> IO (Map PoolKey Int)
poolIndices (Pool IORef PoolState
ref Int
_) [PoolKey]
keys = do
  [(PoolKey, Int)]
ixs <- [PoolKey] -> (PoolKey -> IO (PoolKey, Int)) -> IO [(PoolKey, Int)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [PoolKey]
keys ((PoolKey -> IO (PoolKey, Int)) -> IO [(PoolKey, Int)])
-> (PoolKey -> IO (PoolKey, Int)) -> IO [(PoolKey, Int)]
forall a b. (a -> b) -> a -> b
$ \PoolKey
k -> do
    PoolState
p <- IORef PoolState -> IO PoolState
forall a. IORef a -> IO a
readIORef IORef PoolState
ref
    case PoolKey -> Map PoolKey Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup PoolKey
k (PoolState -> Map PoolKey Int
poolKeys PoolState
p) of
      Just Int
i -> (PoolKey, Int) -> IO (PoolKey, Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PoolKey
k, Int
i)
      Maybe Int
Nothing -> do
        PoolState
p <- PoolState -> IO PoolState
forall {f :: * -> *}.
(PrimState f ~ RealWorld, PrimMonad f) =>
PoolState -> f PoolState
grow PoolState
p
        let i :: Int
i = PoolState -> Int
poolNext PoolState
p
            !c :: Closure
c = PoolKey -> Closure
closureFor PoolKey
k
        MutableArray (PrimState IO) Closure -> Int -> Closure -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray (PoolState -> MutableArray (PrimState IO) Closure
poolArray PoolState
p) Int
i Closure
c
        -- a native run that began before the array last grew holds an old one
        [MutableArray RealWorld Closure]
-> (MutableArray RealWorld Closure -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ (PoolState -> [MutableArray (PrimState IO) Closure]
poolOld PoolState
p) ((MutableArray RealWorld Closure -> IO ()) -> IO ())
-> (MutableArray RealWorld Closure -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \MutableArray RealWorld Closure
old -> Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MutableArray RealWorld Closure -> Int
forall s a. MutableArray s a -> Int
sizeofMutableArray MutableArray RealWorld Closure
old) (MutableArray (PrimState IO) Closure -> Int -> Closure -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray RealWorld Closure
MutableArray (PrimState IO) Closure
old Int
i Closure
c)
        IORef PoolState -> PoolState -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef PoolState
ref PoolState
p {poolNext = i + 1, poolKeys = Map.insert k i (poolKeys p)}
        pure (PoolKey
k, Int
i)
  pure ([(PoolKey, Int)] -> Map PoolKey Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(PoolKey, Int)]
ixs)
  where
    grow :: PoolState -> f PoolState
grow PoolState
p
      | PoolState -> Int
poolNext PoolState
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< MutableArray RealWorld Closure -> Int
forall s a. MutableArray s a -> Int
sizeofMutableArray (PoolState -> MutableArray (PrimState IO) Closure
poolArray PoolState
p) = PoolState -> f PoolState
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PoolState
p
      | Bool
otherwise = do
          let old :: MutableArray (PrimState IO) Closure
old = PoolState -> MutableArray (PrimState IO) Closure
poolArray PoolState
p
              n :: Int
n = MutableArray RealWorld Closure -> Int
forall s a. MutableArray s a -> Int
sizeofMutableArray MutableArray RealWorld Closure
MutableArray (PrimState IO) Closure
old
          MutableArray (PrimState f) Closure
new <- Int -> Closure -> f (MutableArray (PrimState f) Closure)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) (Closure -> f (MutableArray (PrimState f) Closure))
-> Closure -> f (MutableArray (PrimState f) Closure)
forall a b. (a -> b) -> a -> b
$! Closure
natTypeTag
          MutableArray (PrimState f) Closure
-> Int -> MutableArray (PrimState f) Closure -> Int -> Int -> f ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
copyMutableArray MutableArray (PrimState f) Closure
new Int
0 MutableArray (PrimState f) Closure
MutableArray (PrimState IO) Closure
old Int
0 Int
n
          pure PoolState
p {poolArray = new, poolOld = old : poolOld p}

closureFor :: PoolKey -> Closure
closureFor :: PoolKey -> Closure
closureFor = \case
  KeyEnum Reference
r PackedTag
t -> Reference -> PackedTag -> Closure
Enum Reference
r PackedTag
t
  -- the segment's boxes must be evaluated (native code reads through
  -- them); nullSeg's components are CAFs, so force them
  KeyComb CombIx
cix GCombInfo MComb
comb -> let (USeg
u, BSeg
b) = (USeg, BSeg)
nullSeg in USeg
u USeg -> Closure -> Closure
forall a b. a -> b -> b
`seq` BSeg
b BSeg -> Closure -> Closure
forall a b. a -> b -> b
`seq` CombIx -> GCombInfo MComb -> (USeg, BSeg) -> Closure
PAp CombIx
cix GCombInfo MComb
comb (USeg
u, BSeg
b)
  KeyCached CombIx
_ Closure
c -> Closure
c
  KeyRecord RecordShape
s -> RecordShape -> RecordVals -> Closure
RecordG RecordShape
s RecordVals
forall a. Vector a
V.empty
  KeyLit MLit
l -> case MLit
l of
    MT Text
t -> Text -> Closure
Stack.TextC Text
t
    MM Referent
r -> Foreign -> Closure
Foreign (Referent -> Foreign
Stack.WrapReferent Referent
r)
    MY Reference
r -> Foreign -> Closure
Foreign (Reference -> Foreign
Stack.WrapReference Reference
r)
    MLit
_ -> String -> Closure
forall a. HasCallStack => String -> a
error String
"closureFor: unboxed literal"