diff --git a/src/Feynman/Optimization/StateFold.hs b/src/Feynman/Optimization/StateFold.hs
index eb11a058..a8892e61 100644
--- a/src/Feynman/Optimization/StateFold.hs
+++ b/src/Feynman/Optimization/StateFold.hs
@@ -24,6 +24,7 @@ import Data.Map.Strict (Map, (!))
 import qualified Data.Map.Strict as Map
 import Data.Set (Set)
 import qualified Data.Set as Set
+import Control.Monad (msum, liftM)
 import Control.Monad.State.Strict hiding (join)
 import Data.Bits
 import Data.Coerce (coerce)