<commit_before>{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
module Ef.Data.MinHeap where

import Control.Monad
import Control.Monad.ST

import Data.Array
import Data.STRef
import Data.Array.ST



data MinHeap s a
    where

        Heap
            :: STRef s Int
            -> STArray s Int a
            -> MinHeap s a


main = do
    let
        minHeap =
            runST $
                do
                    array <- newListArray (1,1 :: Int) [10 :: Int] :: ST s (STArray s Int Int)
                    (buildMinHeap :: STArray s Int Int -> ST s (MinHeap s Int)) array :: ST s (MinHeap s Int)

        min =
            runST $
                do
                    (extractMinST :: MinHeap s Int -> ST s (Maybe Int)) minHeap

    print min

buildMinHeap
    :: Ord a
    => STArray s Int a
    -> ST s (MinHeap s a)

buildMinHeap array =
    do
        (_,elements) <- getBounds array
        elems <- newSTRef elements
        let
            largestIndex =
                elements `div` 2

            nextLargestIndex =
                pred largestIndex

        forM [largestIndex, nextLargestIndex .. 1] $ \index ->
            minHeapify (Heap elems array) index
        return (Heap elems array)



extractMinST
    :: Ord a
    => MinHeap s a
    -> ST s (Maybe a)

extractMinST (Heap elems heap) =
    do
        elements <- readSTRef elems
        if elements < 1 then
            return Nothing
        else
            do
                min <- readArray heap 1
                newFirst <- readArray heap elements
                writeArray heap 1 newFirst
                writeSTRef elems (pred elements)
                minHeapify (Heap elems heap) 1
                return (Just min)



minHeapify
    :: Ord a
    => MinHeap s a
    -> Int
    -> ST s ()

minHeapify (Heap elems heap) =
    go
    where
        go index =
            do
                elements <- readSTRef elems
                let
                    leftIndex =
                        2 * index

                    rightIndex =
                        leftIndex + 1

                    hasLeft =
                        leftIndex <= elements

                    hasRight =
                        rightIndex <= elements

                    readCurrent =
                        readArray heap index

                    readLeft =
                        readArray heap leftIndex

                    readRight =
                        readArray heap rightIndex

                current <- readCurrent

                small@(smaller,smallerIndex) <-
                    if hasLeft then
                        do
                            left <- readLeft
                            return $
                                if left < current then
                                    (left,leftIndex)
                                else
                                    (current,index)

                    else
                        return (current,index)

                (smallest,indexOfSmallest) <-
                    if hasRight then
                        do
                            right <- readRight
                            return $
                                if right < smaller then
                                    (right,rightIndex)
                                else
                                    small
                    else
                        return small

                when (indexOfSmallest /= index) $
                    do
                        writeArray heap index smallest
                        writeArray heap indexOfSmallest current
                        go indexOfSmallest



-- insert
--     :: Int
--     -> 
<commit_msg>Added failing example.<commit_after>{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE GADTs #-}
-- module Ef.Data.MinHeap where
module Main where

import Control.Monad
import Control.Monad.ST

import Data.Array.ST
import Data.STRef



data MinHeap s a
    where

        Heap
            :: STRef s Int
            -> STArray s Int a
            -> MinHeap s a


main = do
    let
        min =
            runST $
                do
                    array <- newListArray (1,5000000) [1 .. 5000000 :: Int]
                    minHeap <- buildMinHeap array
                    replicateM 10 (extractMinST minHeap)
    print min
    -- print $ head $ reverse [5000000,49999995000,4999 .. 1 .. 1 :: Int]

buildMinHeap
    :: Ord a
    => STArray s Int a
    -> ST s (MinHeap s a)

buildMinHeap array =
    do
        (_,elements) <- getBounds array
        elems <- newSTRef elements
        let
            largestIndex =
                elements `div` 2

            nextLargestIndex =
                pred largestIndex

        forM [largestIndex, nextLargestIndex .. 1] $ \index ->
            minHeapify (Heap elems array) index
        return (Heap elems array)



extractMinST
    :: Ord a
    => MinHeap s a
    -> ST s (Maybe a)

extractMinST (Heap elems heap) =
    do
        elements <- readSTRef elems
        if elements < 1 then
            return Nothing
        else
            do
                min <- readArray heap 1
                newFirst <- readArray heap elements
                writeArray heap 1 newFirst
                writeSTRef elems (pred elements)
                minHeapify (Heap elems heap) 1
                return (Just min)



minHeapify
    :: Ord a
    => MinHeap s a
    -> Int
    -> ST s ()

minHeapify (Heap elems heap) start =
    do
        current <- readArray heap start
        go start current
    where
        {-# INLINE go #-}
        go index current =
            do
                elements <- readSTRef elems
                let
                    leftIndex =
                        2 * index

                    rightIndex =
                        leftIndex + 1

                    hasLeft =
                        leftIndex <= elements

                    hasRight =
                        rightIndex <= elements

                    readLeft =
                        readArray heap leftIndex

                    readRight =
                        readArray heap rightIndex

                small@(smaller,smallerIndex) <-
                    if hasLeft then
                        do
                            left <- readLeft
                            return $
                                if left < current then
                                    (left,leftIndex)
                                else
                                    (current,index)

                    else
                        return (current,index)

                (smallest,indexOfSmallest) <-
                    if hasRight then
                        do
                            right <- readRight
                            return $
                                if right < smaller then
                                    (right,rightIndex)
                                else
                                    small
                    else
                        return small

                when (indexOfSmallest /= index) $
                    do
                        writeArray heap index smallest
                        writeArray heap indexOfSmallest current
                        go indexOfSmallest smallest



-- insert
--     :: Int
--     -> 

{-# INLINE extractMinST #-}
{-# INLINE minHeapify #-}
{-# INLINE buildMinHeap #-}
{'fraction_non_alphanumeric': 0.058258898813491534, 'fraction_numerical': 0.008265564591387815, 'mean_word_length': 0.8839779005524862, 'pattern_counts': {'":': 0, '<': 38, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5288006', 'n_tokens_mistral': 1911, 'n_tokens_neox': 1734, 'n_words': 549}