Variadic.hs

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Data.Kind (Type)
import Foreign.Storable (Storable(..))
import Foreign.Marshal.Array (mallocArray, peekArray, reallocArray, newArray)
import Foreign.Ptr (castPtr, Ptr, plusPtr)
import qualified Foreign.Marshal.Alloc as F (malloc, free)

import Data.List (splitAt)

type family Unfoldr (f :: k1 -> k2 -> k2) (e :: k2) (es :: [k1]) :: k2 where
  Unfoldr f e '[]     = e
  Unfoldr f e (t:ts)  = f t (Unfoldr f e ts)

class Storable a => Alloc (a :: Type) (args :: [Type]) where
  malloc :: Unfoldr (->) (IO (Ptr a)) args

data Stack (a :: Type)
  = Stack
  { elems :: Ptr (Ptr a)
  , size  :: Int
  , off   :: Int
  }

instance Storable a => Storable (Stack a) where
  sizeOf _ = sizeOf @(Ptr (Ptr a)) undefined + sizeOf @Int undefined * 2
  alignment _ = 8
  peek ptr = do
    let sizeofPtr = sizeOf @(Ptr (Ptr a)) undefined
        sizeofInt = sizeOf @Int undefined
    Stack <$> (peekByteOff ptr 0 :: IO (Ptr (Ptr a)))
          <*> (peekByteOff ptr sizeofPtr :: IO Int)
          <*> (peekByteOff ptr (sizeofPtr + sizeofInt) :: IO Int)

  poke ptr Stack{..} = do
    let sizeofPtr = sizeOf @(Ptr (Ptr a)) undefined
        sizeofInt = sizeOf @Int undefined
    pokeByteOff ptr 0 elems
    pokeByteOff ptr sizeofPtr size
    pokeByteOff ptr (sizeofPtr + sizeofInt) off

instance Storable a => Alloc (Stack a) '[Int] where
  malloc initSize = do
    stack <- F.malloc @(Stack a)
    elems <- mallocArray @(Ptr a) initSize
    let size = initSize
        off  = 0
    poke stack Stack{..}
    return stack

push :: Storable a => Ptr (Stack a) -> Ptr a -> IO ()
push ptr elem = do
  st@Stack{..} <- peek ptr
  st@Stack{..} <- if off + 1 == size
                  then let newSize = size * 2
                       in Stack <$> (castPtr <$> reallocArray elems newSize)
                                <*> pure newSize
                                <*> pure off
                  else pure st
  list         <- newArray =<< replaceAt off elem <$> peekArray size elems
  poke ptr (Stack list size (off + 1))

pop :: Storable a => Ptr (Stack a) -> IO a
pop ptr = do
  st@Stack{..} <- peek ptr
  elem         <- peek =<< (!! off) <$> peekArray size elems
  poke ptr (Stack elems size (off - 1))
  return elem

freeStack :: Storable a => Ptr (Stack a) -> IO ()
freeStack ptr = do
  st@Stack{..} <- peek ptr
  mapM_ F.free =<< peekArray size elems
  F.free elems
  F.free ptr

replaceAt :: Int -> a -> [a] -> [a]
replaceAt n newElem list =
  let (begin, _ : end) = splitAt n list
  in begin <> (newElem : end)

main :: IO ()
main = do
  stackPtr <- malloc @(Stack Int) @'[Int] 5
  n        <- F.malloc @Int
  poke n 3

  push stackPtr n
  print =<< pop stackPtr

  F.free n
  freeStack stackPtr


---------------------------------------------------------------------------------------------

{-

While this definitely typechecks (`malloc @(Stack Int) @'[Int] :: Int -> IO (Ptr (Stack Int))`), the compiled program
segfaults almost immediately on the `Storable` instance of `Int`. I believe this is a problem with my hand-made `Storable`
instances, but I don't know.
Anyway, I proved my point that we can create typesafe variadic functions in Haskell with some type level trickery.

-}