WIP: An alternative implementation of thunks

hnix.cabal

@@ -467,6 +467,9 @@ library
       Nix.Thunk
       Nix.Thunk.Basic
       Nix.Thunk.Standard
+      Nix.Thunk.StableId
+      Nix.Thunk.FreshStableIdT
+      Nix.Thunk.Separate
       Nix.Type.Assumption
       Nix.Type.Env
       Nix.Type.Infer
@@ -496,6 +499,7 @@ library
     , exceptions
     , filepath
     , free
+    , ghc-prim
     , hashing
     , hnix-store-core
     , http-client

src/Nix/Exec.hs

@@ -589,6 +589,9 @@ newtype Lazy t (f :: * -> *) m a = Lazy
         , MonadReader (Context (Lazy t f m) t)
         )
 
+mapLazy :: (forall x. m x -> m x) -> Lazy t f m a -> Lazy t f m a
+mapLazy f (Lazy a) = Lazy $ mapReaderT (mapStateT f) a
+
 instance MonadTrans (Lazy t f) where
   lift = Lazy . lift . lift
 

src/Nix/Thunk/FreshStableIdT.hs

@@ -0,0 +1,50 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Nix.Thunk.FreshStableIdT (FreshStableIdT, runFreshStableIdT) where
+
+import Nix.Thunk
+import Nix.Thunk.StableId
+import Control.Monad.Reader
+import Control.Monad.State.Strict
+import Control.Monad.Ref
+import Control.Monad.Catch
+import Control.Applicative
+#ifdef MIN_VERSION_haskeline
+import System.Console.Haskeline.MonadException (MonadException)
+#endif
+
+newtype FreshStableIdT m a = FreshStableIdT (ReaderT StableId (StateT Int m) a)
+  deriving
+    ( Functor
+    , Applicative
+    , Monad
+    , MonadRef
+    , MonadAtomicRef
+    , MonadCatch
+    , MonadThrow
+    , MonadIO
+    , MonadFix
+    , MonadPlus
+    , Alternative
+#ifdef MIN_VERSION_haskeline
+    , MonadException
+#endif
+    )
+
+instance MonadTrans FreshStableIdT where
+  lift = FreshStableIdT . lift . lift
+
+runFreshStableIdT :: Monad m => StableId -> FreshStableIdT m a -> m a
+runFreshStableIdT root (FreshStableIdT a) = evalStateT (runReaderT a root) 0
+
+instance Monad m => MonadThunkId (FreshStableIdT m) where
+  type ThunkId (FreshStableIdT m) = StableId
+  freshId = FreshStableIdT $ do
+    root <- ask
+    n <- get
+    put $ succ n
+    return $ cons n root

src/Nix/Thunk/Separate.hs

@@ -0,0 +1,208 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE DeriveFoldable #-}
+{-# LANGUAGE DeriveFunctor #-}
+{-# LANGUAGE DeriveTraversable #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE FunctionalDependencies #-}
+{-# LANGUAGE RankNTypes #-}
+
+{-# OPTIONS_GHC -Wno-orphans #-}
+
+module Nix.Thunk.Separate
+  ( NThunkF (..)
+  , MonadSeparateThunk
+  , SeparateThunkT
+  , runSeparateThunkT
+  , askThunkCache
+  , buildThunk
+  , separateThunkId
+  , queryValue
+  , queryThunk
+  , forceThunk
+  , forceEffects
+  , valueRef
+  , thunkValue
+  , newThunkCache
+  , Forcer (..)
+  ) where
+
+import Control.Exception hiding (catch)
+import Control.Applicative
+import Control.Monad.Catch
+import Control.Monad.Reader
+import Control.Monad.Ref
+import Data.Map.Strict (Map)
+import qualified Data.Map.Strict as Map
+
+import Nix.Thunk
+import Nix.Thunk.StableId
+import Nix.Thunk.FreshStableIdT
+
+-- | The type of very basic thunks
+data NThunkF forcer v
+    = Value v
+    | Thunk (ThunkId forcer) (forcer v)
+
+instance (Eq v, Eq (ThunkId forcer)) => Eq (NThunkF forcer v) where
+    Value x == Value y = x == y
+    Thunk x _ == Thunk y _ = x == y
+    _ == _ = False              -- jww (2019-03-16): not accurate...
+
+instance Show v => Show (NThunkF forcer v) where
+    show (Value v) = show v
+    show (Thunk _ _) = "<thunk>"
+
+type MonadSeparateThunk m = (MonadAtomicRef m) --TODO: ThunkId allocation also needs to be sufficiently deterministic
+
+newtype ThunkCache m v = ThunkCache (Ref m (Map StableId (Maybe v)))
+
+class Monad forcer => Forcer forcer v m | forcer -> v, forcer -> m where
+  liftSeparateThunkT :: SeparateThunkT forcer v m a -> forcer a
+  mapSeparateThunkT :: (forall x. SeparateThunkT forcer v m x -> SeparateThunkT forcer v m x) -> forcer a -> forcer a
+
+--TODO: HashMap?
+newtype SeparateThunkT (forcer :: * -> *) v m a = SeparateThunkT (ReaderT (ThunkCache m v) (FreshStableIdT m) a)
+  deriving
+    ( Functor
+    , Applicative
+    , Monad
+    , MonadRef
+    , MonadAtomicRef
+    , MonadCatch
+    , MonadThrow
+    , MonadIO
+    , MonadPlus
+    , Alternative
+    , MonadFix
+    )
+
+newThunkCache :: (MonadRef m, Ref m ~ Ref m') => m (ThunkCache m' v)
+newThunkCache = ThunkCache <$> newRef Map.empty
+
+askThunkCache :: Monad m => SeparateThunkT forcer v m (ThunkCache m v)
+askThunkCache = SeparateThunkT ask
+
+runSeparateThunkT :: Monad m => StableId -> ThunkCache m v -> SeparateThunkT forcer v m a -> m a
+runSeparateThunkT root c (SeparateThunkT a) = runFreshStableIdT root $ runReaderT a c
+
+instance MonadTrans (SeparateThunkT forcer v) where
+    lift = SeparateThunkT . lift . lift
+
+instance Monad m => MonadThunkId (SeparateThunkT forcer v m) where
+    type ThunkId (SeparateThunkT forcer v m) = StableId
+    freshId = SeparateThunkT $ lift freshId
+
+{-
+instance (MonadSeparateThunk m, MonadCatch m)
+  => MonadThunk (NThunkF forcer v) (SeparateThunkT forcer v m) v where
+    thunk     = buildThunk
+    thunkId   = \case
+        Value _   -> Nothing
+        Thunk n _ -> Just n
+    query     = queryValue
+    queryM    = queryThunk
+    force     = forceThunk
+    forceEff  = forceEffects
+    wrapValue = valueRef
+    getValue  = thunkValue
+-}
+
+separateThunkId :: NThunkF forcer v -> Maybe (ThunkId forcer)
+separateThunkId = \case
+  Value _   -> Nothing
+  Thunk n _ -> Just n
+
+valueRef :: v -> NThunkF forcer v
+valueRef = Value
+
+thunkValue :: NThunkF forcer v -> Maybe v
+thunkValue (Value v) = Just v
+thunkValue _ = Nothing
+
+buildThunk :: (Monad m, Forcer forcer v m, ThunkId forcer ~ StableId) => forcer v -> SeparateThunkT forcer v m (NThunkF forcer v)
+buildThunk action = do
+    freshThunkId <- freshId
+    c <- askThunkCache
+    return $ Thunk freshThunkId $ mapSeparateThunkT (lift . runSeparateThunkT freshThunkId c) action
+
+queryValue :: NThunkF forcer v -> a -> (v -> a) -> a
+queryValue (Value v) _ k = k v
+queryValue _ n _ = n
+
+queryThunk :: (Forcer forcer v m, Monad m, MonadAtomicRef m, ThunkId forcer ~ StableId) => NThunkF forcer v -> forcer a -> (v -> forcer a) -> forcer a
+queryThunk (Value v) _ k = k v
+queryThunk (Thunk tid _) n k = do
+    ThunkCache c <- liftSeparateThunkT $ SeparateThunkT ask
+    mOldVal <- liftSeparateThunkT $ atomicModifyRef' c $ \old ->
+        -- Try to insert Nothing into the given key, but if something is already
+        -- there, just leave it
+        let (mOldVal, !new) = Map.insertLookupWithKey (\_ _ oldVal -> oldVal) tid Nothing old
+        in (new, mOldVal)
+    case mOldVal of
+        Nothing -> do
+            result <- n -- Not computed, inactive
+            -- This is the only case where we've actually changed c, so restore it
+            liftSeparateThunkT $ atomicModifyRef' c $ \old -> (Map.delete tid old, ())
+            return result
+        Just Nothing -> n -- Active
+        Just (Just v) -> k v -- Computed, inactive
+
+forceThunk
+    :: forall m v a forcer.
+    ( MonadAtomicRef m
+    , MonadThrow forcer
+    , MonadCatch forcer
+    , Forcer forcer v m
+    , ThunkId forcer ~ StableId
+    )
+    => NThunkF forcer v -> (v -> forcer a) -> forcer a
+forceThunk (Value v) k = k v
+forceThunk (Thunk tid action) k = do
+    ThunkCache c <- liftSeparateThunkT $ SeparateThunkT ask
+    mOldVal <- liftSeparateThunkT $ atomicModifyRef' c $ \old ->
+        -- Try to insert Nothing into the given key, but if something is already
+        -- there, just leave it
+        let (mOldVal, !new) = Map.insertLookupWithKey (\_ _ oldVal -> oldVal) tid Nothing old
+        in (new, mOldVal)
+    case mOldVal of
+        Nothing -> do -- Not computed, inactive
+            v <- catch action $ \(e :: SomeException) -> do
+                -- This is the only case where we've actually changed c, so restore it
+                _ <- liftSeparateThunkT $ atomicModifyRef' c $ \old -> (Map.delete tid old, ())
+                throwM e
+            liftSeparateThunkT $ atomicModifyRef' c $ \old -> (Map.insert tid (Just v) old, ())
+            k v
+        Just Nothing -> throwM $ ThunkLoop $ show tid
+        Just (Just v) -> k v -- Computed, inactive
+
+forceEffects
+  :: ( MonadAtomicRef m
+     , Forcer forcer v m
+     , ThunkId forcer ~ StableId
+     )
+  => NThunkF forcer v -> (v -> forcer r) -> forcer r
+forceEffects (Value v) k = k v
+forceEffects (Thunk tid action) k = do
+    ThunkCache c <- liftSeparateThunkT $ SeparateThunkT ask
+    mOldVal <- liftSeparateThunkT $ atomicModifyRef' c $ \old ->
+        -- Try to insert Nothing into the given key, but if something is already
+        -- there, just leave it
+        let (mOldVal, !new) = Map.insertLookupWithKey (\_ _ oldVal -> oldVal) tid Nothing old
+        in (new, mOldVal)
+    case mOldVal of
+        Nothing -> do -- Not computed, inactive
+            v <- action
+            liftSeparateThunkT $ atomicModifyRef' c $ \old -> (Map.insert tid (Just v) old, ())
+            k v
+        Just Nothing -> return $ error "Loop detected"
+        Just (Just v) -> k v -- Computed, inactive

src/Nix/Thunk/StableId.hs

@@ -0,0 +1,88 @@
+{-# LANGUAGE MultiWayIf #-}
+{-# LANGUAGE MagicHash #-}
+
+-- Equivalent to [Int], but with near-O(1) amortized comparison
+module Nix.Thunk.StableId (StableId, nil, cons, uncons) where
+
+import Data.IORef
+import System.IO.Unsafe
+import GHC.Prim
+import Data.Hashable
+import Data.List (unfoldr)
+import Data.Ord
+
+import Debug.Trace
+
+--TODO: If we have a really long chain, we will keep leaking memory; what can we do about this?
+
+data StableId = StableId
+  { _stableId_value :: {-# UNPACK #-} !Int
+  , _stableId_hash :: {-# UNPACK #-} !Int
+  , _stableId_parent :: {-# UNPACK #-} !(IORef StableId)
+  }
+
+{-# NOINLINE nil #-} -- If nil is not a single value on the heap, infinite recursion can result
+nil :: StableId
+nil = StableId 0 0 $ unsafePerformIO $ newIORef $ error "nil"
+
+cons :: Int -> StableId -> StableId
+cons v p@(StableId _ ph _) = StableId v (hash (v, ph)) $ unsafeDupablePerformIO $ newIORef p
+
+uncons :: StableId -> Maybe (Int, StableId)
+uncons s = if _stableId_parent s == _stableId_parent nil
+  then Nothing
+  else Just
+  ( _stableId_value s
+  , unsafeDupablePerformIO $ readIORef $ _stableId_parent s
+  )
+
+--TODO: Reimplement Eq in terms of Ord?
+instance Eq StableId where
+  a == b = trace (show a <> " == " <> show b) $ if
+    | _stableId_parent a == _stableId_parent b -- We're the exact same heap object
+      -> True
+    | _stableId_hash a /= _stableId_hash b || _stableId_value a /= _stableId_value b -- We're definitely different
+      -> False
+    | _stableId_parent a == _stableId_parent nil || _stableId_parent b == _stableId_parent nil -- One of them is nil, but the other isn't.  Note that this relies on nil being unique. --TODO: Can we avoid this?
+      -> False
+    | otherwise -- Different objects, but same value and hash.  These are either the same value or a hash collision.
+      -> unsafeDupablePerformIO $ do
+           pa <- readIORef $ _stableId_parent a
+           pb <- readIORef $ _stableId_parent b
+           case reallyUnsafePtrEquality# pa pb of
+             -- Parents are different objects
+             0# -> if pa == pb
+               then do writeIORef (_stableId_parent b) pa -- Parents are equivalent, so unify
+                       return True
+               else return False -- Parents are not equivalent, so leave them alone
+             -- Parents are the same object already
+             _ -> return True
+
+instance Ord StableId where
+  a `compare` b = trace (show a <> " `compare` " <> show b) $ case comparing _stableId_hash a b <> comparing _stableId_value a b of
+    LT -> LT
+    GT -> GT
+    EQ -> case _stableId_parent a == _stableId_parent b of
+      True -> EQ
+      False ->
+        if _stableId_parent a == _stableId_parent nil then LT else --TODO: Can we avoid this?
+        if _stableId_parent b == _stableId_parent nil then GT else
+        unsafeDupablePerformIO $ do
+          pa <- readIORef $ _stableId_parent a
+          pb <- readIORef $ _stableId_parent b
+          case reallyUnsafePtrEquality# pa pb of
+            -- Parents are different objects
+            0# -> case pa `compare` pb of
+              LT -> return LT
+              GT -> return GT
+              EQ -> do
+                writeIORef (_stableId_parent b) pa
+                return EQ
+            -- Parents are the same object already
+            _ -> return EQ
+
+toList :: StableId -> [Int]
+toList = unfoldr uncons
+
+instance Show StableId where
+  showsPrec n = showsPrec n . toList