Thunk.Basic: refactoring the thunk lock/unlock procedure

Author: Anton-Latukha

src/Nix/Thunk/Basic.hs

@@ -32,11 +32,25 @@ import           Nix.Utils                      ( bool
 data Deferred m v = Computed v | Deferred (m v)
   deriving (Functor, Foldable, Traversable)
 
+-- | It is a reference (@ref-tf: Ref m@), and as such also holds @Bool@ lock.
+type ThunkRef m = (Var m Bool)
+
+-- | Reference (@ref-tf: Ref m v@) to a value that thunk holds.
+type ThunkValueRef m v = Var m (Deferred m v)
+
+-- | @ref-tf@ lock instruction for @Ref m@ (@ThunkRef@).
+lock :: Bool -> (Bool, Bool)
+lock = (True, )
+
+-- | @ref-tf@ unlock instruction for @Ref m@ (@ThunkRef@).
+unlock :: Bool -> (Bool, Bool)
+unlock = (False, )
+
 -- * Data type for thunks: @NThunkF@
 
 -- | The type of very basic thunks
 data NThunkF m v
-  = Thunk (ThunkId m) (Var m Bool) (Var m (Deferred m v))
+  = Thunk (ThunkId m) (ThunkRef m) (ThunkValueRef m v)
 
 instance (Eq v, Eq (ThunkId m)) => Eq (NThunkF m v) where
   Thunk x _ _ == Thunk y _ _ = x == y
@@ -52,16 +66,19 @@ type MonadBasicThunk m = (MonadThunkId m, MonadVar m)
 instance (MonadBasicThunk m, MonadCatch m)
   => MonadThunk (NThunkF m v) m v where
 
+  -- | Return thunk ID
   thunkId :: NThunkF m v -> ThunkId m
   thunkId (Thunk n _ _) = n
 
+  -- | Create new thunk
   thunk :: m v -> m (NThunkF m v)
   thunk action =
     do
       freshThunkId <- freshId
       Thunk freshThunkId <$> newVar False <*> newVar (Deferred action)
 
-  -- | Non-blocking query
+  -- | Non-blocking query, return value if @Computed@,
+  -- return first argument otherwise.
   queryM :: m v -> NThunkF m v -> m v
   queryM n (Thunk _ _ ref) =
     do
@@ -88,33 +105,34 @@ instance (MonadBasicThunk m, MonadCatch m)
 
 -- *** United body of `force*`
 
+-- | If @m v@ is @Computed@ - returns is
 forceMain
   :: ( MonadBasicThunk m
     , MonadCatch m
     )
   => NThunkF m v
   -> m v
-forceMain (Thunk n active ref) =
+forceMain (Thunk n thunkRef thunkValRef) =
   do
     deferred
       pure
       (\ action ->
         do
-          lockThunk <- atomicModifyVar active (True, )
+          lockedIt <- atomicModifyVar thunkRef lock
           bool
             (throwM $ ThunkLoop $ show n)
             (do
               v <- catch action $ \(e :: SomeException) ->
                 do
-                  _ <- atomicModifyVar active (False, )
+                  _unlockedIt <- atomicModifyVar thunkRef unlock
                   throwM e
-              writeVar ref (Computed v)
-              _unlockThunk <- atomicModifyVar active (False, )
+              writeVar thunkValRef (Computed v)
+              _unlockedIt <- atomicModifyVar thunkRef unlock
               pure v
             )
-            (not lockThunk)
+            (not lockedIt)
       )
-      =<< readVar ref
+      =<< readVar thunkValRef
 {-# inline forceMain #-} -- it is big function, but internal, and look at its use.
 
 
@@ -129,22 +147,23 @@ instance (MonadBasicThunk m, MonadCatch m)
     -> m r
     -> NThunkF m v
     -> m r
-  queryMF k n (Thunk _ active ref) =
+  queryMF k n (Thunk _ thunkRef thunkValRef) =
     do
-      thunkIsAvaliable <- not <$> atomicModifyVar active (True, )
+      lockedIt <- atomicModifyVar thunkRef (True, )
       bool
         n
         go
-        thunkIsAvaliable
+        (not lockedIt)
     where
       go =
         do
-          eres <- readVar ref
+          eres <- readVar thunkValRef
           res  <-
-            case eres of
-              Computed v   -> k v
-              Deferred _mv -> n
-          _ <- atomicModifyVar active (False, )
+            deferred
+              k
+              (const n)
+              eres
+          _unlockedIt <- atomicModifyVar thunkRef (False, )
           pure res
 
   forceF
@@ -165,10 +184,12 @@ instance (MonadBasicThunk m, MonadCatch m)
     -> m (NThunkF m v)
   furtherF k t@(Thunk _ _ ref) =
     do
-      _ <- atomicModifyVar ref $
-        \x -> case x of
-          Computed _ -> (x, x)
-          Deferred d -> (Deferred (k d), x)
+      _modifiedIt <- atomicModifyVar ref $
+        \x ->
+          deferred
+            (const (x, x))
+            (\ d -> (Deferred (k d), x))
+            x
       pure t