Search code examples
haskellspace-leak

STUArray s i e - space leak only when i == Int?


I am confused by the behaviour of the following snipped:

import Data.Int
import Data.Array.ST
import Control.Monad.ST

{-# INLINE fib #-}
fib _ 0 = return 0
fib _ 1 = return 1
fib c n = do
  f1 <- memo c (fib c) (n-1)
  f2 <- memo c (fib c) (n-2)
  return (f1+f2)

newtype C a = C a

{-# INLINE memo #-}
memo (C a) f k = do
  e <- readArray a k
  if e == (-1)
     then do
       v <- f k
       writeArray a k v
       return v
     else return e

evalFib :: Int -> Int
evalFib n = runST $ do
  a <- newArray (0,n) (-1) :: ST s (STUArray s Int Int)
  fib (C a) n

main = print $ evalFib 120000

When compiled with -O2 it stack-overflows (showing 20M of memory in use). The confusing part is that it actually works as expected (no stack overflow and 9M memory in use) if any of the following changes is made:

  • Int64 is used instead of Int: (giving evalFib :: Int64 -> Int and STUArray s Int64 Int). In fact any Int* (Int32, Int16, etc) will do the trick as well as Word or Word*;
  • newtype C a is removed from the picture;
  • data C a = C !a is used instead of newtype

I am trying to get my head around this behaviour: is it a bug in GHC/array module (it shows identical behaviour in 7.4.2 and 7.6.2) or is it supposed to work that way?

PS The funny thing is that when I try to compile it with ghc array.hs -O2 -fext-core to see the differences in core produced both GHC versions fail with "ghc: panic! (the 'impossible' happened)". No luck here either..


Solution

  • Looking at the core from 7.6.1, with -O2 and -dsuppress-uniques, the function that does the work, Main.main_$spoly_$wa is structurally (almost) identical whether I use int or Int64 as the index type. Since the core is long and complicated, here's the diff output:

    $ diff Int_core.dump-simpl Int64_core.dump-simpl 
    11,12c11,12
    <           (Data.Array.Base.STUArray s GHC.Types.Int GHC.Types.Int)
    <           (Main.C (Data.Array.Base.STUArray s GHC.Types.Int GHC.Types.Int))
    ---
    >           (Data.Array.Base.STUArray s GHC.Int.Int64 GHC.Types.Int)
    >           (Main.C (Data.Array.Base.STUArray s GHC.Int.Int64 GHC.Types.Int))
    26,27c26,27
    <             (Data.Array.Base.STUArray s GHC.Types.Int GHC.Types.Int)
    <             (Main.C (Data.Array.Base.STUArray s GHC.Types.Int GHC.Types.Int)))
    ---
    >             (Data.Array.Base.STUArray s GHC.Int.Int64 GHC.Types.Int)
    >             (Main.C (Data.Array.Base.STUArray s GHC.Int.Int64 GHC.Types.Int)))
    

    Different index types, these are of course different.

    33,40d32
    <           l :: GHC.Types.Int
    <           [LclId]
    <           l = GHC.Types.I# sc } in
    <         let {
    <           u :: GHC.Types.Int
    <           [LclId]
    <           u = GHC.Types.I# sc1 } in
    <         let {
    

    For index type Int, GHC produces somewhat more informative errors for out-of-bounds indices, for that it needs the lower and upper bound of the valid indices. (The default implementation of index is not overridden in the instance Ix Int64.)

    45,46c37
    <           GHC.Types.False ->
    <             case poly_$w$j5 (GHC.Types.I# a) l u of wild2 { };
    ---
    >           GHC.Types.False -> case GHC.Arr.hopelessIndexError of wild1 { };
    

    Different error, indexError vs. hopelessIndexError. The following differences also only concern index errors.

    49,50c40
    <               GHC.Types.False ->
    <                 case poly_$w$j5 (GHC.Types.I# a) l u of wild2 { };
    ---
    >               GHC.Types.False -> case GHC.Arr.hopelessIndexError of wild2 { };
    58c48
    <                     case poly_$w$j4 y (GHC.Types.I# sc2) of wild3 { };
    ---
    >                     case poly_$w$j3 y (GHC.Types.I# sc2) of wild4 { };
    62c52
    <                         case poly_$w$j4 y (GHC.Types.I# sc2) of wild5 { };
    ---
    >                         case poly_$w$j3 y (GHC.Types.I# sc2) of wild5 { };
    77,78c67
    <                                 GHC.Types.False ->
    <                                   case poly_$w$j3 (GHC.Types.I# a1) l u of wild6 { };
    ---
    >                                 GHC.Types.False -> case GHC.Arr.hopelessIndexError of wild6 { };
    81,82c70
    <                                     GHC.Types.False ->
    <                                       case poly_$w$j3 (GHC.Types.I# a1) l u of wild7 { };
    ---
    >                                     GHC.Types.False -> case GHC.Arr.hopelessIndexError of wild7 { };
    

    Now once more the different index type:

    110c98
    <                                                                       GHC.Types.Int
    ---
    >                                                                       GHC.Int.Int64
    152c140
    <                                                 s GHC.Types.Int GHC.Types.Int>)>)
    ---
    >                                                 s GHC.Int.Int64 GHC.Types.Int>)>)
    

    And finally, 0 and 1 have gotten different top-level names.

    177,178c165,166
    <       0 -> (# sc5, lvl5 #);
    <       1 -> (# sc5, lvl6 #)
    ---
    >       0 -> (# sc5, lvl #);
    >       1 -> (# sc5, lvl1 #)
    

    So the entire code that does the actual work is identical. Yet the one causes a stack overflow (though only just, -K9M suffices [-K8731K is enough here, -K8730K not]) and the other doesn't.

    The difference is indeed caused by the index errors. The code with Int indices allocates two boxed Ints in every recursive call that the Int64 code doesn't allocate, because

    Main.main_$spoly_$wa [Occ=LoopBreaker]
      :: forall s.
         GHC.Prim.Int#
         -> GHC.Prim.Int#
         -> GHC.Prim.Int#
         -> GHC.Prim.MutableByteArray# s
         -> (GHC.Prim.~#)
              *
              (Data.Array.Base.STUArray s GHC.Types.Int GHC.Types.Int)
              (Main.C (Data.Array.Base.STUArray s GHC.Types.Int GHC.Types.Int))
         -> GHC.Prim.Int#
         -> GHC.Prim.State# s
         -> (# GHC.Prim.State# s, GHC.Types.Int #)
    

    carries around two references to the array.

    That uses more stack, and these boxed Ints have to be garbage collected, which causes the much larger GC figures. Additionally, the thunk for the index error is a bit larger than the hopelessIndexError thunk.

    Now, if you help the compiler by

    • removing the newtype wrapper
    • making the function strict in the array (by bang patterns or data C a = C !a)

    or some other ways, it produces better code that manages without a stack overflow for the given argument, since there is only one reference to the array in the worker, and thus the allocation of the boxed Ints for the bounds isn't needed.

    Note that this algorithm causes a stack overflow for slightly larger arguments even with the help for the compiler.