feat(cdcl): add activity-ordered variable heap

src/Sift/Cdcl/Heap.hs

@@ -0,0 +1,120 @@
+{-# LANGUAGE FlexibleContexts #-}
+module Sift.Cdcl.Heap
+  ( VarHeap
+  , newHeap
+  , heapInsert
+  , heapIncrease
+  , heapRemoveMax
+  , heapInHeap
+  , heapIsEmpty
+  ) where
+
+import Control.Monad (when)
+import Control.Monad.ST (ST)
+import Data.STRef
+import Data.Array.ST
+
+data VarHeap s = VarHeap
+  { hArr  :: STUArray s Int Int
+  , hPos  :: STUArray s Int Int
+  , hSize :: STRef s Int
+  , hAct  :: STUArray s Int Double
+  }
+
+newHeap :: Int -> STUArray s Int Double -> ST s (VarHeap s)
+newHeap n act = do
+  arr  <- newArray (0, n - 1) 0
+  pos  <- newArray (0, n - 1) (-1)
+  size <- newSTRef 0
+  pure (VarHeap arr pos size act)
+
+heapInHeap :: VarHeap s -> Int -> ST s Bool
+heapInHeap h v = do
+  p <- readArray (hPos h) v
+  pure (p >= 0)
+
+heapIsEmpty :: VarHeap s -> ST s Bool
+heapIsEmpty h = do
+  n <- readSTRef (hSize h)
+  pure (n == 0)
+
+heapInsert :: VarHeap s -> Int -> ST s ()
+heapInsert h v = do
+  present <- heapInHeap h v
+  when (not present) $ do
+    n <- readSTRef (hSize h)
+    writeArray (hArr h) n v
+    writeArray (hPos h) v n
+    writeSTRef (hSize h) (n + 1)
+    siftUp h n
+
+heapIncrease :: VarHeap s -> Int -> ST s ()
+heapIncrease h v = do
+  p <- readArray (hPos h) v
+  when (p >= 0) $ siftUp h p
+
+heapRemoveMax :: VarHeap s -> ST s Int
+heapRemoveMax h = do
+  n    <- readSTRef (hSize h)
+  top  <- readArray (hArr h) 0
+  lastV <- readArray (hArr h) (n - 1)
+  writeArray (hArr h) 0 lastV
+  writeArray (hPos h) lastV 0
+  writeArray (hPos h) top (-1)
+  writeSTRef (hSize h) (n - 1)
+  when (n - 1 > 0) $ siftDown h 0
+  pure top
+
+siftUp :: VarHeap s -> Int -> ST s ()
+siftUp h start = do
+  x  <- readArray (hArr h) start
+  ax <- readArray (hAct h) x
+  let go j
+        | j <= 0    = finish j
+        | otherwise = do
+            let parent = (j - 1) `div` 2
+            yp <- readArray (hArr h) parent
+            ay <- readArray (hAct h) yp
+            if ax > ay
+              then do
+                writeArray (hArr h) j yp
+                writeArray (hPos h) yp j
+                go parent
+              else finish j
+      finish j = do
+        writeArray (hArr h) j x
+        writeArray (hPos h) x j
+  go start
+
+siftDown :: VarHeap s -> Int -> ST s ()
+siftDown h start = do
+  n  <- readSTRef (hSize h)
+  x  <- readArray (hArr h) start
+  ax <- readArray (hAct h) x
+  let go j = do
+        let l = 2 * j + 1
+        if l >= n
+          then finish j
+          else do
+            let r = l + 1
+            cl  <- readArray (hArr h) l
+            acl <- readArray (hAct h) cl
+            (childIdx, childVar, childAct) <-
+              if r < n
+                then do
+                  cr  <- readArray (hArr h) r
+                  acr <- readArray (hAct h) cr
+                  if acr > acl
+                    then pure (r, cr, acr)
+                    else pure (l, cl, acl)
+                else pure (l, cl, acl)
+            if childAct > ax
+              then do
+                writeArray (hArr h) j childVar
+                writeArray (hPos h) childVar j
+                go childIdx
+              else finish j
+      finish j = do
+        writeArray (hArr h) j x
+        writeArray (hPos h) x j
+  go start