Consider avoiding reader for SetM

[meooow25]

containers/containers/src/Data/Graph.hs

Line 584 in 0d85628

newtype SetM s a = SetM { runSetM :: STUArray s Vertex Bool -> ST s a }

We currently use a reader for the GHC version of SetM, which means the recursive go function passes around the visited array all the time. This generates Core like

letrec {
  $sgo_sbDD [Occ=LoopBreaker,
             Dmd=SC(S,C(1,C(1,C(1,C(1,C(1,!P(L,L)))))))]
    :: Int#
       -> Int#
       -> Int#
       -> MutableByteArray# RealWorld
       -> [Int]
       -> State# RealWorld
       -> (# State# RealWorld, [Tree Vertex] #)
  [LclId[StrictWorker([~, ~, ~, ~, !])],
   Arity=6,
   Str=<L><L><L><L><1L><L>,
   Unf=OtherCon []]
  $sgo_sbDD
    = \ ...

GHC unboxes the array, which means four static arguments.

If we wrote dfs manually, we would not likely not write it in this manner. We would create the array once, declare the recursive go with it in scope, and simply refer to it.

dfs g vs = runST $ do
  m <- newArray ...
  let go = ...
           ... use m ...
  go vs

So why don't we abstract that?

-run          :: Bounds -> (forall s. SetM s a) -> a
-contains     :: Vertex -> SetM s Bool
-include      :: Vertex -> SetM s ()
...
+run :: Bounds -> (forall s. (Int -> ST s Bool) -> (Int -> ST s ()) -> ST s a) -> a
+run bnds f = runST $ do
+  m <- newArray bnds False
+  f (readArray m) (\v -> writeArray m v True)

 dfs :: Graph -> [Vertex] -> [Tree Vertex]
-dfs g vs0 = run (bounds g) $ go vs0
-  where
+dfs g vs0 = run (bounds g) $ \contains include ->
+  let
     go [] = pure []
...

Now our go is simply in ST and does not pass around the array.

The IntSet based implementation for non-GHC remains largely unchanged.

Effect on performance: Nothing significant, at least on my machine. A few benchmarks seem to be faster by 5-10%. So this is mostly motivated by me wanting to improve what it compiles to, but I will share proper numbers once the changes are ready for a PR.