1 files changed, 6 insertions, 14 deletions
diff --git a/test/Main.hs b/test/Main.hs
index 83271fc..1ad7f75 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -14,7 +14,6 @@ import Control.Monad.Trans.Class (lift)
 import Control.Monad.IO.Class (liftIO)
 import Control.Monad.Trans.State
 import Data.Bifunctor
--- import qualified Data.Functor.Product as Product
 import Data.Int (Int64)
 import Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
@@ -22,8 +21,7 @@ import qualified Data.Text as T
 import Hedgehog
 import qualified Hedgehog.Gen as Gen
 import qualified Hedgehog.Range as Range
-import Test.Tasty
-import Test.Tasty.Hedgehog
+import Test.Framework
 
 import Array
 import AST
@@ -209,12 +207,10 @@ adTestGen name expr envGenerator =
   let env = knownEnv @env
       exprS = simplifyFix expr
   in
-  withCompiled env expr $ \getprimalfun ->
-  withCompiled env (simplifyFix expr) $ \getprimalSfun ->
+  withCompiled env expr $ \primalfun ->
+  withCompiled env (simplifyFix expr) $ \primalSfun ->
   testGroup name
     [testProperty "compile primal" $ property $ do
-       primalfun <- liftIO getprimalfun
-       primalSfun <- liftIO getprimalSfun
        input <- forAllWith (showEnv env) envGenerator
 
        let outPrimalI = interpretOpen False input expr
@@ -225,9 +221,8 @@ adTestGen name expr envGenerator =
        outPrimalSC <- liftIO $ primalSfun input
        diff outPrimalSI (closeIsh' 1e-8) outPrimalSC
 
-    ,withCompiled (dne env) (dfwdDN exprS) $ \getdnfun ->
+    ,withCompiled (dne env) (dfwdDN exprS) $ \dnfun ->
      testProperty "compile fwdAD" $ property $ do
-       dnfun <- liftIO getdnfun
        input <- forAllWith (showEnv env) envGenerator
        dinput <- forAllWith (showEnv (dne env)) $ extendDNE env input
        let (outDNI1, outDNI2) = interpretOpen False dinput (dfwdDN expr)
@@ -235,11 +230,8 @@ adTestGen name expr envGenerator =
        diff outDNI1 (closeIsh' 1e-8) outDNC1
        diff outDNI2 (closeIsh' 1e-8) outDNC2
 
-    ,withResource (makeFwdADArtifactCompile env exprS) (\_ -> pure ()) $ \getfwdartifactC ->
+    ,withResource (makeFwdADArtifactCompile env exprS) (\_ -> pure ()) $ \fwdartifactC ->
      testProperty "chad" $ property $ do
-       primalSfun <- liftIO getprimalSfun
-       fwdartifactC <- liftIO getfwdartifactC
-
        annotate (concat (unSList (\t -> ppSTy 0 t ++ " -> ") env) ++ ppSTy 0 (typeOf expr))
 
        let dtermChad0 = ELet ext (EConst ext STF64 1.0) $ chad' defaultConfig env expr
@@ -288,7 +280,7 @@ adTestGen name expr envGenerator =
     envScalars SNil SNil = []
     envScalars (t `SCons` ts) (Value x `SCons` xs) = tanScalars t x ++ envScalars ts xs
 
-withCompiled :: SList STy env -> Ex env t -> (IO (SList Value env -> IO (Rep t)) -> TestTree) -> TestTree
+withCompiled :: SList STy env -> Ex env t -> ((SList Value env -> IO (Rep t)) -> TestTree) -> TestTree
 withCompiled env expr = withResource (compile env expr) (\_ -> pure ())
 
 term_build1_sum :: Ex '[TArr N1 (TScal TF64)] (TScal TF64)