diff --git a/Setup.hs b/Setup.hs
--- a/Setup.hs
+++ b/Setup.hs
@@ -1,2 +1,6 @@
-import Distribution.Simple
-main = defaultMain
+module Main where
+
+import Distribution.Extra.Doctest (defaultMainWithDoctests)
+
+main :: IO ()
+main = defaultMainWithDoctests "doctests"
diff --git a/pcg-random.cabal b/pcg-random.cabal
--- a/pcg-random.cabal
+++ b/pcg-random.cabal
@@ -1,5 +1,5 @@
 name:                pcg-random
-version:             0.1.3.5
+version:             0.1.3.6
 synopsis:            Haskell bindings to the PCG random number generator.
 description:
   PCG is a family of simple fast space-efficient statistically good
@@ -25,9 +25,15 @@
 Bug-reports:         http://github.com/cchalmers/pcg-random/issues
 copyright:           (c) 2014-2015. Christopher Chalmers <c.chalmers@me.com>
 category:            System
-build-type:          Simple
+build-type:          Custom
 cabal-version:       >=1.10
 
+custom-setup
+ setup-depends:
+   base >= 4 && <5,
+   Cabal,
+   cabal-doctest >= 1 && <1.1
+
 source-repository head
   type:     git
   location: git://github.com/cchalmers/pcg-random.git
@@ -46,7 +52,7 @@
   other-extensions:    BangPatterns, CPP, ForeignFunctionInterface
   build-depends:
     base        >=4.3 && < 5,
-    primitive   >=0.4 && < 0.7,
+    primitive   >=0.4 && < 0.8,
     random      >=1.0 && < 2.0,
     bytestring,
     entropy
@@ -75,4 +81,5 @@
   hs-source-dirs:   test
   build-depends:
     base,
-    doctest
+    doctest,
+    pcg-random
diff --git a/src/System/Random/PCG/Fast/Pure.hs b/src/System/Random/PCG/Fast/Pure.hs
--- a/src/System/Random/PCG/Fast/Pure.hs
+++ b/src/System/Random/PCG/Fast/Pure.hs
@@ -168,7 +168,7 @@
 -- | Generate a new seed using single 'Word64'.
 --
 --   >>> initFrozen 0
---   FrozenGen 1
+--   F 1
 initFrozen :: Word64 -> FrozenGen
 initFrozen w = F (w .|. 1)
 
@@ -183,7 +183,7 @@
 -- | Initialize a generator a single word.
 --
 --   >>> initialize 0 >>= save
---   FrozenGen 1
+--   F 1
 initialize :: PrimMonad m => Word64 -> m (Gen (PrimState m))
 initialize a = restore (initFrozen a)
 
diff --git a/test/doctest.hs b/test/doctest.hs
--- a/test/doctest.hs
+++ b/test/doctest.hs
@@ -1,12 +1,7 @@
-import Test.DocTest
+module Main where
 
-main = doctest
-  [ "src/System/Random/PCG.hs"
-  , "src/System/Random/PCG/Class.hs"
-  , "src/System/Random/PCG/Pure.hs"
-  , "src/System/Random/PCG/Fast.hs"
-  , "src/System/Random/PCG/Single.hs"
-  , "dist/build/c/pcg-advance-64.o"
-  , "dist/build/c/pcg-output-64.o"
-  , "dist/build/c/pcg-rngs-64.o"
-  ]
+import Build_doctests (flags, pkgs, module_sources)
+import Test.DocTest (doctest)
+
+main :: IO ()
+main = doctest $ flags ++ pkgs ++ module_sources