An indexed monad to count mallocs and frees

PseudoLinearTypes.hs

{-# LANGUAGE
    DataKinds,
    FlexibleInstances,
    FlexibleContexts,
    PolyKinds,
    ScopedTypeVariables,
    TypeApplications,
    TypeFamilies,
    TypeInType,
    TypeOperators #-}

import Data.Coerce (coerce)
import GHC.TypeLits

-- | An IO computation with starting memory state i and final state j.
--
-- The memory state is a fresh counter and a set of live pointers.
newtype LIO (i :: Ix) (j :: Ix) a = LIO (IO a)

-- | Starting with no pointers, we must end with no pointers left.
runLIO :: LIO '(0, '[]) '(j, '[]) () -> IO ()
runLIO (LIO io) = io

-- * Composition

(>>=.) :: forall i j k a b. LIO i j a -> (a -> LIO j k b) -> LIO i k b
(>>=.) = coerce ((>>=) @IO @a @b)

pure_ :: forall i a. a -> LIO i i a
pure_ = coerce (pure @IO @a)

-- * Memory state

type Id = Nat  -- Unique type-level identifier

type Ix = (Id, [Id])

newtype Ptr (x :: Nat) = Ptr Int

-- * malloc, adding a new element to the state

type family Incr (i :: Ix) :: Ix
type instance Incr '(x, xs) = '(x + 1, x ': xs)

type family Fst (i :: Ix) :: Id
type instance Fst '(x, _) = x

malloc :: LIO i (Incr i) (Ptr (Fst i))
malloc = LIO (putStrLn "Allocated stuff!" >> pure (Ptr 0)) -- dummy

-- * free, removing an element from the state

type family Rmv (y :: Id) (i :: Ix) :: Ix
type instance Rmv y '(x, xs) = '(x, RmvList y xs)

type family RmvList (y :: Id) (xs :: [Id]) :: [Id] where
  RmvList y (y ': xs) = xs
  RmvList y (x ': xs) = x ': RmvList y xs

free :: Ptr x -> LIO i (Rmv x i) ()
free _ = LIO (putStrLn "Freed stuff!") -- dummy

-- Forgetting to free, or freeing twice the same pointer, causes a type error.
main :: IO ()
main = runLIO $
  malloc >>=. \p1 ->
  malloc >>=. \p2 ->
  free p2 >>=. \() ->
  free p1