Quicksort with local mutation in ST monad with Java collections in Frege

QuicksortInversion.fr

module helloworld.QuicksortInversion where
 
--Quicksort with inversion
qsort :: Ord a => [a] -> ([a], Int)
qsort xs = ST.run go where
  go = do
    jlist <- ArrayList.fromList xs
    inv <- Ref.new 0 -- initialize the inversion as 0
    let qsortImperative !from !to
            | from < to = do
                Ref.modify (+ (to - from)) inv -- update inversion
                pivotPos <- medianPos jlist from to -- median as pivot
                pivot <- jlist.unsafeGet pivotPos
                swap jlist pivotPos from
                !newPivotPos <- qpartition jlist from to pivot
                swap jlist from newPivotPos
                qsortImperative from (newPivotPos- 1)
                qsortImperative (newPivotPos + 1) to
            | otherwise = return ()
    size <- jlist.size
    qsortImperative 0 (size - 1)
    sorted <- jlist.toList
    invCount <- Ref.get inv
    return (sorted, invCount)

qpartition :: Ord a => Mutable s (JList a) -> Int -> Int -> a -> ST s Int 
qpartition xs !from !to pivot = go (from + 1) (from + 1) where
    go !i !j
        | j <= to = do
                elem <- xs.unsafeGet j
                case compare elem pivot of
                  LT        -> swap xs j i >> go (i + 1)(j + 1)
                  otherwise -> go i (j + 1)
        | otherwise = return (i - 1)
 
medianOn :: Ord b => (a -> b) -> a -> a -> a -> a
medianOn f a b c =
    let (large, small) = if f a > f b then (a, b) else (b, a) in
    if f c > f large then large else (if f c < f small then small else c)
   
medianPos :: Ord a => Mutable s (JList a) ->Int -> Int -> ST s Int
medianPos xs from to = do
  let mid = (from + to) `quot` 2
  felem <- xs.unsafeGet from
  telem <- xs.unsafeGet to
  melem <- xs.unsafeGet mid
  return . snd $ medianOn fst (felem, from)(telem, to) (melem, mid)
 
--Java Definitions
data IndexOutOfBoundsException = native java.lang.IndexOutOfBoundsException
derive Exceptional IndexOutOfBoundsException
 
data Collection a = native java.util.Collection
 
data Iterator e = native java.util.Iterator where
  native hasNext :: Mutable s (Iterator e) -> ST s Bool
  native next :: Mutable s (Iterator e) -> ST s e
  native remove :: Mutable s (Iterator e) -> ST s ()
 
data JList a = native java.util.List where
  native add :: Mutable s (JList a) -> a -> ST s ()
  native get :: Mutable s (JList a) -> Int -> ST s (Maybe a)
    throws IndexOutOfBoundsException
  native unsafeGet get :: Mutable s (JList a) -> Int -> ST s a
    throws IndexOutOfBoundsException --might return null
  native set :: Mutable s (JList a) -> Int -> a -> ST s a
  native size :: Mutable s (JList a) -> ST s Int
  native isEmpty :: Mutable s (JList a) -> ST s Bool
  native iterator :: Mutable s (JList a) -> STMutable s (Iterator a)
   
  toList :: Mutable s (JList a) -> ST s [a]
  toList xs = do
    itr <- xs.iterator
    let loop acc = do
          hasNext <- itr.hasNext
          if hasNext
              then do
                next <- itr.next
                loop (next:acc)
              else return (reverse acc)
    loop []
   
data ArrayList a = native java.util.ArrayList where
  native new :: () -> STMutable s (ArrayList a)
                | Mutable s (Collection a) -> STMutable s (ArrayList a)
  fromList :: [a] -> STMutable s (ArrayList a)
  fromList xs = do
    jlist <- ArrayList.new ()
    let loop [] = return jlist
        loop (x:xs) = jlist.add x >> loop xs
    loop xs
   
native swap java.util.Collections.swap ::Mutable s (JList a) -> Int -> Int -> ST s ()
 
main _ = println $ qsort [2, 4, 3, 1, 5, 8, 6]