diff --git a/benchmarks/Main.hs b/benchmarks/Main.hs
--- a/benchmarks/Main.hs
+++ b/benchmarks/Main.hs
@@ -21,7 +21,7 @@
     ]
   where
     mkBenchmark :: Int -> Benchmark
-    mkBenchmark size = env (pure [0 .. size]) $ \lst -> bench (show size) $ nfAppIO recursiveSize lst
+    mkBenchmark size = env (pure [0 .. size]) $ \lst -> bench (show size) $ nfAppIO (runHeapsize . recursiveSize) lst
 
 byteArrayBenchmark :: Benchmark
 byteArrayBenchmark = bgroup "ByteArray#" $
@@ -31,14 +31,14 @@
     mkBenchmark :: Int -> Benchmark
     mkBenchmark size = env
       (newByteArray size)
-      (\lst -> bench (show size) $ nfAppIO recursiveSize lst)
+      (\lst -> bench (show size) $ nfAppIO (runHeapsize . recursiveSize) lst)
 
 
 data C = C Int (IORef C)
   deriving (Generic, NFData)
 
 circularBenchmark :: Benchmark
-circularBenchmark = env mkT (\c -> bench "Circular" (nfAppIO recursiveSize c))
+circularBenchmark = env mkT (\c -> bench "Circular" (nfAppIO (runHeapsize . recursiveSize) c))
   where
     mkT :: IO C
     mkT = do
diff --git a/heapsize.cabal b/heapsize.cabal
--- a/heapsize.cabal
+++ b/heapsize.cabal
@@ -1,13 +1,13 @@
 cabal-version:      3.0
 name:               heapsize
-version:            0.1
+version:            0.2.0
 license:            BSD-3-Clause
 license-file:       LICENSE
 category:           GHC, Debug, Development
 build-type:         Simple
 author:             Michail Pardalos <mpardalos@gmail.com>
-maintainer:         Michail Pardalos <mpardalos@gmail.com>
-copyright:          Michail Pardalos 2020
+maintainer:         Pepe Iborra <pepeiborra@gmail.com>
+copyright:          Michail Pardalos 2020, Pepe Iborra 2020-2021
 synopsis:           Determine the size of runtime data structures
 description:        heapsize is a tool to determine the size data structures.
                     Determining the size of recursive data structures is
@@ -20,11 +20,14 @@
     cbits/heapsize_prim.cmm
   Default-Language: Haskell2010
   Build-depends:
-      base >= 4.12 && < 4.15
-    , deepseq >= 1.3 && < 1.5
+      base >= 4.12 && < 5.0
+    , deepseq >= 1.3
+    , exceptions
     , ghc-heap >= 8.6
-    , unordered-containers ==0.2.*
-    , hashable ==1.3.*
+    , hashtables
+    , transformers
+    , unordered-containers >=0.2
+    , hashable >=1.3
     , primitive
   Hs-source-dirs: src/
   Ghc-options: -Wall
@@ -44,4 +47,4 @@
 
 source-repository head
     type:     git
-    location: https://github.com/mpardalos/heapsize.git
+    location: https://github.com/pepeiborra/heapsize.git
diff --git a/src/HeapSize.hs b/src/HeapSize.hs
--- a/src/HeapSize.hs
+++ b/src/HeapSize.hs
@@ -1,7 +1,8 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE UnboxedTuples #-}
-{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
@@ -18,29 +19,41 @@
  -}
 module HeapSize (
   recursiveSize,
-  recursiveSizeNoGC,
   recursiveSizeNF,
-  closureSize
+  closureSize,
+  Heapsize,
+  runHeapsize
   )
   where
 
 import Control.DeepSeq (NFData, force)
+import Control.Exception (throwIO)
+import Control.Monad
+import Control.Monad.Catch
+import Control.Monad.IO.Class
+import Control.Monad.Trans.Maybe
+import Control.Monad.Trans.Reader
+import Data.IORef
+import Data.Hashable
+import qualified Data.HashTable.IO as HT
+import Data.Maybe (isJust, isNothing)
+import Data.Typeable (Typeable)
 
 import GHC.Exts hiding (closureSize#)
 import GHC.Arr
 import GHC.Exts.Heap hiding (size)
-import qualified Data.HashSet as H
-import Data.IORef
-import Data.Hashable
-
-import Control.Monad
+import qualified Data.Foldable as F
 
 import System.Mem
+import System.Mem.Weak
+import Debug.Trace
 
 foreign import prim "aToWordzh" aToWord# :: Any -> Word#
 foreign import prim "unpackClosurePtrs" unpackClosurePtrs# :: Any -> Array# b
 foreign import prim "closureSize" closureSize# :: Any -> Int#
 
+----------------------------------------------------------------------------
+
 -- | Get the *non-recursive* size of an closure in words
 closureSize :: a -> IO Int
 closureSize x = return (I# (closureSize# (unsafeCoerce# x)))
@@ -51,6 +64,54 @@
       let nelems = I# (sizeofArray# pointers)
       in pure (fmap Box $ Array 0 (nelems - 1) nelems pointers)
 
+--------------------------------------------------------------------------------
+
+data HeapsizeState = HeapsizeState
+  {
+    -- | A mutable seen set
+    closuresSeen :: HT.BasicHashTable HashableBox (),
+    -- | A counter for the seen set
+    seenSizeRef  :: IORef Int,
+    -- | Did the GC run since the computation start?
+    gcDetect     :: GcDetector
+  }
+
+-- | A one-shot device for detecting garbage collections
+newtype GcDetector = GcDetector {gcSinceCreation :: IO Bool}
+
+gcDetector :: IO GcDetector
+gcDetector = do
+  ref <- newIORef ()
+  w <- mkWeakIORef ref (return ())
+  return $ GcDetector $ isNothing <$> deRefWeak w
+
+newtype Heapsize a = Heapsize
+  { _unHeapsize :: ReaderT HeapsizeState (MaybeT IO) a}
+  deriving (Applicative, Functor, Monad, MonadIO, MonadCatch, MonadMask, MonadThrow)
+
+initSize :: Int
+initSize = 10000000
+
+--   A garbage collection is performed before the size is calculated, because
+--   the garbage collector would make heap walks difficult.
+--   Returns `Nothing` if the count is interrupted by a garbage collection
+runHeapsize :: Heapsize a -> IO (Maybe a)
+runHeapsize (Heapsize comp) = do
+
+  -- initialize the mutable state
+  !closuresSeen <- HT.newSized initSize
+  seenSizeRef   <- newIORef 0
+
+  -- Perform a major GC
+  performMajorGC
+
+  -- Create a GC detector for the duration of the entire computation
+  gcDetect <- gcDetector
+
+  runMaybeT $ runReaderT comp HeapsizeState{..}
+
+--------------------------------------------------------------------------------
+
 -- | Calculate the recursive size of GHC objects in Bytes. Note that the actual
 --   size in memory is calculated, so shared values are only counted once.
 --
@@ -67,45 +128,54 @@
 --   ($!! from Control.DeepSeq) to force full evaluation before calculating the
 --   size.
 --
---   A garbage collection is performed before the size is calculated, because
---   the garbage collector would make heap walks difficult.
---
 --   This function works very quickly on small data structures, but can be slow
 --   on large and complex ones. If speed is an issue it's probably possible to
 --   get the exact size of a small portion of the data structure and then
 --   estimate the total size from that.
-recursiveSize :: a -> IO Int
-recursiveSize x = performGC >> recursiveSizeNoGC x
+recursiveSize :: a -> Heapsize Int
+recursiveSize x = Heapsize $ do
+  HeapsizeState{..} <- ask
+  accSizeRef <- liftIO $ newIORef 0
+  let checkGC = gcSinceCreation gcDetect >>= \abort -> when abort $ throwIO Interrupted
+  let
+    go :: [ Box ] -> IO ()
+    go [] = return ()
+    go (b@(Box y) : rest) = do
+      let addr = W# (aToWord# y)
+      !seen <- isJust <$> HT.lookup closuresSeen (HashableBox b)
 
--- | Same as `recursiveSize` except without performing garbage collection first.
---   Useful if you want to measure the size of many objects in sequence. You can
---   call `performGC` once at first and then use this function to avoid multiple
---   unnecessary garbage collections.
-recursiveSizeNoGC :: a -> IO Int
-recursiveSizeNoGC x = do
-  state <- newIORef (0, H.empty)
-  go state (asBox x)
+      next <- if seen || isBadAddress addr then return [] else do
+          -- always check that GC has not happened before deref pointers
+          checkGC
+          thisSize <- closureSize y
+          checkGC
+          next <- getClosures y
 
-  fst <$> readIORef state
-  where
-    go :: IORef (Int, H.HashSet HashableBox) -> Box -> IO ()
-    go state b@(Box y) = do
-      (_, closuresSeen) <- readIORef state
+          HT.insert closuresSeen (HashableBox b) ()
+          modifyIORef' accSizeRef (+ thisSize)
+          modifyIORef' seenSizeRef succ
+          return (F.toList next)
+      go (next ++ rest)
 
-      when (not $ H.member (HashableBox b) closuresSeen) $ do
-        thisSize <- closureSize y
-        modifyIORef state $ \(size, _) ->
-          (size + thisSize, H.insert (HashableBox b) closuresSeen)
+  liftIO (go [asBox x]) `catch` \Interrupted -> do
+    seen <- liftIO $ readIORef seenSizeRef
+    liftIO $ traceIO ("SEEN: " <> show seen)
+    mzero
+  liftIO (readIORef accSizeRef)
 
-        mapM_ (go state) =<< getClosures y
+data Interrupted = Interrupted deriving (Show, Typeable)
+instance Exception Interrupted
 
 -- | Calculate the recursive size of GHC objects in Bytes after calling
 -- Control.DeepSeq.force on the data structure to force it into Normal Form.
 -- Using this function requires that the data structure has an `NFData`
 -- typeclass instance.
-
-recursiveSizeNF :: NFData a => a -> IO Int
+recursiveSizeNF :: NFData a => a -> Heapsize Int
 recursiveSizeNF = recursiveSize . force
+
+isBadAddress :: Word -> Bool
+isBadAddress 0 = True
+isBadAddress _ = False
 
 newtype HashableBox = HashableBox Box
     deriving newtype Show