STUArray woes

So a little while ago, I was trying to write a nice, clear version of the language shootout n-body problem, using the ST monad, and STUArrays. Though I ran into a snag; I was getting segmentation faults, for very small values of n (< 23), and not knowing enough about such things, I have no idea where to begin in diagnosing the problem. So, I shall post my code here, and hope that someone might have some ideas. I've talked to Don Stewart about it, but he hasn't got back to me yet.

(All the MUArray instance stuff is from Don's work, which I nicked)

So, here we have it:

> {-# OPTIONS -O2 -funbox-strict-fields #-}
> {-# LANGUAGE BangPatterns          #-}
> {-# LANGUAGE UnboxedTuples         #-}
> {-# LANGUAGE MagicHash             #-}
> {-# LANGUAGE MultiParamTypeClasses #-}
> 
> module Main where
> 
> import System.Environment
> import Data.Array.Base
> import GHC.Exts
> import GHC.ST
> import GHC.Arr
> import Control.Monad
> import Control.Monad.ST
> import Data.List
> import Foreign.Storable
> import Text.Printf
> import Data.STRef
> 
> 
> data Vec3 = V !Double !Double !Double deriving (Show)
> data Body = B {  
>                  pos ::  {-# UNPACK #-} !Vec3,
>                  vel ::  {-# UNPACK #-} !Vec3,
>                  mass :: !Double
>                } deriving (Show)
> 
> 
> main = do
>     n <- getArgs >>= \xs -> if null xs then return 1000 else (readIO.head) xs
>     let (e,e') = runST $ do
>             ps <- planets
>             offset ps
>             e1 <- energy ps
>             advn n ps
>             e2 <- energy ps
>             return (e1,e2)
>     prtnum e
>     prtnum e'
> 
> 
> offset ps = do
>     v <- newSTRef (V 0 0 0)
>     (B p ve m) <- readArray ps 0
>     forM_ [0..numPs] $ \i -> do
>         p <- readArray ps i
>         modifySTRef v (.+. momentum p)
>     modifySTRef v (./ solar_mass)
>     v' <- readSTRef v
>     writeArray ps 0 (B p (ve.-. v') m)
> 
>     
> 
> advn 0 ps = do
>     return ps
> advn n ps = do
>     advance deltat ps
>     -- e <- energy ps
>     -- unsafeIOToST (prtnum e)
>     advn (n-1) ps
> 
> advance dt ps = do
>     forM_ [0..numPs] $ \i -> do
>         (B p v m) <- readArray ps i
>         vvar <- newSTRef v
>         forM_ [i+1..numPs] $ \j -> do
>             (B p' v' m') <- readArray ps j
>             let dp = p .-. p'
>                 dst = mag dp
>                 magnit = dt / (dst * dst * dst)
>                 change = magnit *. dp
>                 sca1 = m' *. change
>                 sca2 = m *. change
>             -- writeArray ps I (B p (v .-. sca1) m)
>             modifySTRef vvar (.-. sca1)
>             writeArray ps j (B p' (v' .+. sca2) m')
>         vvar' <- readSTRef vvar
>         writeArray ps I (move dt (B p vvar' m))
> 
> 
> energy ps = do
>     e <- newSTRef (0 :: Double)
>     
>     forM_ [0..numPs] $ \i -> do
>         b1 <- readArray ps i
>         modifySTRef e ((+) (kinetic b)1)
>         
>         forM_ [i+1..numPs] $ \j -> do
>             b2 <- readArray ps j
>             let x = potential b1 b2
>             modifySTRef e (subtract x)
>     e' <- readSTRef e
>     return e'
> 
> 
> 
> move dt (B p v m) = B (p .+. dt *. v) v m
> 
> 
> 
> 
> -----------------------
> solar_mass, deltat, days_per_year :: Double
> days_per_year = 365.24
> dp = days_per_year
> solar_mass    = 4 * pi * pi
> deltat       = 0.01
> numPs       = length planets' - 1
> 
> planets'' = newArray_ (0,numPs) :: ST s (STUArray s Int Body)
>     -- (B (V 0 0 0) (V 0 0 0) 1)
> 
> planets = do
>     arr <- planets''
>     forM_ (zip [0..] planets') $ \(i,p) -> do
>         writeArray arr I p
>     return arr
> 
> planets' =
>     [B  (V  0.0
>             0.0
>             0.0) 
>         (V  0.0
>             0.0
>             0.0) 
>         solar_mass,
>         
>      B  (V  4.84143144246472090e+00
>             (-1.16032004402742839e+00)
>             (-1.03622044471123109e-01)) 
>         (V  (1.66007664274403694e-03*dp) 
>             (7.69901118419740425e-03*dp)
>             ((-6.90460016972063023e-05)*dp)) 
>         (9.54791938424326609e-04 * solar_mass),
>         
>      B  (V  8.34336671824457987e+00
>             4.12479856412430479e+00
>             (-4.03523417114321381e-01))
>         (V  ((-2.76742510726862411e-03)*dp)
>             (4.99852801234917238e-03*dp)
>             (2.30417297573763929e-05*dp)) 
>         (2.85885980666130812e-04 * solar_mass),
>         
>      B  (V  1.28943695621391310e+01
>             (-1.51111514016986312e+01)
>             (-2.23307578892655734e-01)) 
>         (V  (2.96460137564761618e-03*dp)
>             (2.37847173959480950e-03*dp)
>             ((-2.96589568540237556e-05)*dp)) 
>         (4.36624404335156298e-05 * solar_mass),
>         
>      B  (V  1.53796971148509165e+01
>             (-2.59193146099879641e+01)
>             1.79258772950371181e-01) 
>         (V  (2.68067772490389322e-03*dp)
>             (1.62824170038242295e-03*dp)
>             ((-9.51592254519715870e-05)*dp)) 
>         (5.15138902046611451e-05 * solar_mass)
>     ]
> 
> 
> -----------------------
> 
> infix 4 .-.
> infix 4 .+.
> -- infix 5 .*.
> infixr 5 *.
> infixl 5 ./
> 
> {-# INLINE (.+.) #-}
> {-# INLINE (.-.) #-}
> -- {-# INLINE (.*.) #-}
> V x y z .+. V x' y' z' = V (x+x') (y+y') (z+z')  
> V x y z .-. V x' y' z' = V (x-x') (y-y') (z-z')  
> -- V x y z .*. V x' y' z' = V (x*x') (y*y') (z*z')
> 
> {-# INLINE (*.) #-}
> n *. (V x y z) = V (n*x) (n*y) (n*z)
> {-# INLINE (./) #-}
> V x y z ./ n = V (x/n) (y/n) (z/n)
> 
> {-# INLINE momentum #-}
> momentum :: Body -> Vec3
> momentum (B _ v m) = m *. v
> 
> {-# INLINE kinetic #-}
> kinetic x =  0.5 * mass x * (mag2.vel) x
> 
> 
> potential a b =
>     let dst = dist' a b
>     in if dst <= 0
>         then 0
>         else (mass a * mass b) / dist' a b
> 
> {-# INLINE mag #-}
> mag :: Vec3 -> Double
> mag = sqrt . mag2
> 
> {-# INLINE mag2 #-}
> mag2 (V x y z) = x*x + y*y + z*z
> 
> {-# INLINE dist #-}
> dist v1 v2 = sqrt $ dist2 v1 v2
> 
> {-# INLINE dist2 #-}
> dist2 v1 v2 = mag2 $ v1 .-. v2
> 
> {-# INLINE dist' #-}
> dist' b1 b2 = dist (pos b1) (pos b2)
> 
> prtnum x = putStrLn $ printf "%.9f" x
> 
> --------------------
> bodyScale :: Int# -> Int#
> bodyScale n = (scale *# 7#) *# n
>      where I# scale = sizeOf (undefined :: Double)
> 
> instance MArray (STUArray s) Body (ST s) where
> 
>     {-# INLINE getBounds #-}
>     getBounds (STUArray l h _ _) = return (l,h)
> 
>     {-# INLINE getNumElements #-}
>     getNumElements (STUArray _ _ n _) = return n
> 
>     {-# INLINE unsafeNewArray_ #-}
>     unsafeNewArray_ (l,u) = unsafeNewArraySTUArray_ (l,u) bodyScale -- bytes 
> 
>     {-# INLINE newArray_ #-}
>     newArray_ bounds = newArray bounds (B (V 0 0 0) (V 0 0 0) 1)
>     
> 
>     {-# INLINE unsafeRead #-}
>     unsafeRead (STUArray _ _ _ arr) (I# i) = ST $ \s1 ->
>         let j = bodyScale i in
>         case readDoubleArray# arr j         s1 of { (# s2, xd    #) ->
>         case readDoubleArray# arr (j +# 1#) s2 of { (# s3, yd    #) ->
>         case readDoubleArray# arr (j +# 2#) s3 of { (# s4, zd    #) ->
>         case readDoubleArray# arr (j +# 3#) s4 of { (# s5, vxd   #) ->
>         case readDoubleArray# arr (j +# 4#) s5 of { (# s6, vyd   #) ->
>         case readDoubleArray# arr (j +# 5#) s6 of { (# s7, vzd   #) ->
>         case readDoubleArray# arr (j +# 6#) s7 of { (# s8, massd #) ->
>         (# s8, B { pos    = V (D# xd) (D# yd) (D# zd)
>                     , vel    = V (D# vxd) (D# vyd) (D# vzd)
>                     , mass = D# massd } #) }}}}}}}
> 
>     {-# INLINE unsafeWrite #-}
>     unsafeWrite (STUArray _ _ _ arr) (I# i)
>                 (B (V (D# xd) (D# yd) (D# zd)) (V (D# vxd) (D# vyd) (D# vzd)) (D# massd))
>         = ST $ \s1 ->
>         let j = bodyScale i in
>         case writeDoubleArray# arr j          xd   s1  of {s2 ->
>         case writeDoubleArray# arr (j +# 1#)  yd   s2  of {s3 ->
>         case writeDoubleArray# arr (j +# 2#)  zd   s3  of {s4 ->
>         case writeDoubleArray# arr (j +# 3#) vxd   s4  of {s5 ->
>         case writeDoubleArray# arr (j +# 4#) vyd   s5  of {s6 ->
>         case writeDoubleArray# arr (j +# 5#) vzd   s6  of {s7 ->
>         case writeDoubleArray# arr (j +# 6#) massd s7  of {s8 ->
>         (# s8, () #) }}}}}}}
>