module AntiMemo ( AntiMemo, antimemoize, readAntiMemo,
                  (|++|),
                ) where

import Control.Monad ( liftM2 )
#ifdef ANTIMEMOIZE
import System.IO.Unsafe ( unsafePerformIO )
import Data.IORef ( newIORef, writeIORef, readIORef, IORef )
import System.Mem.Weak ( mkWeakPtr, deRefWeak, Weak )
#include "impossible.h"
#endif

antimemoize :: (a -> b) -> a -> AntiMemo b
readAntiMemo :: AntiMemo b -> b

(|++|) :: AntiMemo [a] -> AntiMemo [a] -> AntiMemo [a]
(|++|) = liftM2 (++)

instance Ord a => Ord (AntiMemo a) where
    x <= y = (readAntiMemo x) <= (readAntiMemo y)

#ifdef ANTIMEMOIZE
data AntiMemo a = AntiMemo (IO a) (IORef (Weak a))

instance Monad AntiMemo where
    thea@(AntiMemo ioa _) >>= f =
        case f (readAntiMemo thea) of
        (AntiMemo _ iowb) -> AntiMemo (unantimemoize f `fmap` ioa) iowb
    return x = unsafePerformIO $ do wb <- mkWeakPtr x Nothing
                                    iowb <- newIORef wb
                                    return (AntiMemo (return x) iowb)

instance Functor AntiMemo where
    fmap f thea@(AntiMemo ioa _) =
        case readAntiMemo thea of
        a -> AntiMemo (f `fmap` ioa) $
             unsafePerformIO $ mkWeakPtr (f a) Nothing >>= newIORef

instance Eq a => Eq (AntiMemo a) where
    x@(AntiMemo _ ix) == y@(AntiMemo _ iy)
        = ix == iy || (readAntiMemo x) == (readAntiMemo y)

unantimemoize :: (a -> AntiMemo b) -> (a -> b)
unantimemoize f a = readAntiMemo $ f a

antimemoize f a = unsafePerformIO $ do wb <- mkWeakPtr (f a) Nothing
                                       iowb <- newIORef wb
                                       return (AntiMemo (return $ f a) iowb)

readAntiMemo (AntiMemo ioa iowa) = unsafePerformIO $
    do ma <- readIORef iowa >>= deRefWeak
       case ma of Nothing -> do a <- ioa
                                wa <- mkWeakPtr a Nothing
                                writeIORef iowa wa
                                return a
                  Just a -> return a

#else

newtype AntiMemo a = AM a

instance Functor AntiMemo where
    fmap f (AM a) = AM (f a)

instance Monad AntiMemo where
    (AM a) >>= f = f a
    return x = AM x

instance Eq a => Eq (AntiMemo a) where
    x == y = readAntiMemo x == readAntiMemo y

readAntiMemo (AM x) = x
antimemoize f a = AM (f a)

#endif