如何修复 Haskell 中的内存泄漏（thunk 泄漏？）？

Question

我正在尝试修复此代码中的内存泄漏。它不会在小输入十分钟内终止，并且在使用 stack exec 运行时会出现“内存不足”错误。我认为内存泄漏是在newMap的构建中。我认为这可能是一个 thunk 泄漏，因为它在递归函数内部。我尝试使用 Strict GHC 扩展并使用 seq 和 deepseq 以及以下代码的变体。

我认为泄漏在这里：

    !existingPolysets <- get
    let newMap = (List.foldl' (\accumulator coefficient ->
                  let newMonomial = (Monomial coefficient degree)
                      innerMap = addOneMonomialToPolySets newMonomial (PolynomialConstraints x y bound maxDegree) existingPolysets accumulator in
                    Map.union innerMap accumulator) Map.empty [-bound..bound])
    ($!!) modify' (const newMap)

下面是剩下的代码：

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE BangPatterns #-}
module Dag where

import Control.Monad.State.Strict
import Control.Monad (liftM, ap)
import qualified Data.List as List
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import GHC.Generics (Generic)
import Data.Hashable
import Control.DeepSeq

data PolynomialConstraints = PolynomialConstraints Integer Integer Integer Integer
-- x y bound degree
  deriving (Show, Eq, Ord, Generic)

data Index = Index Integer Integer -- degree, output
  deriving (Show, Eq, Ord, Generic)
type MyState = Map.Map Index PolynomialSetWithIdentifier

instance NFData Index
instance NFData Monomial
instance NFData Identifier
instance NFData SubPolynomialSetWithIdentifier
instance NFData PolynomialSetWithIdentifier
instance NFData PolynomialConstraints

instance Hashable Monomial
instance Hashable Identifier
instance Hashable SubPolynomialSetWithIdentifier
instance Hashable PolynomialSetWithIdentifier

data PolynomialSetWithIdentifier = PolynomialSetWithIdentifier [SubPolynomialSetWithIdentifier] Identifier
  deriving (Show, Eq, Ord, Generic)
data SubPolynomialSetWithIdentifier = SubPolynomialSetWithIdentifier Monomial PolynomialSetWithIdentifier Identifier
  deriving (Show, Eq, Ord, Generic)

data Identifier = Identifier Int
  deriving (Show, Eq, Ord, Generic)

type PolynomialsSolutions a = State MyState a

data Monomial = Monomial Integer Integer
-- coeff, degree
  deriving (Show, Eq, Ord, Generic)

emptyPolynomialSetWithID = (PolynomialSetWithIdentifier [] (Identifier $ hash $ getSubPolynomialSetsIdentifiers []))

getSubPolynomialSetsIdentifiers :: [SubPolynomialSetWithIdentifier] -> [Identifier] -- is this good or should it be [Int] ? more typing seems safer
-- returns a list the identifiers of the SubPolynomialSets in subPolynomialSetsWithID
getSubPolynomialSetsIdentifiers = List.map (\(SubPolynomialSetWithIdentifier monomial polynomialSetWithID identifier) -> identifier)

getPolynomialSetIdentifier :: PolynomialSetWithIdentifier -> Identifier
-- returns the identifier of a PolynomialSet
getPolynomialSetIdentifier (PolynomialSetWithIdentifier subPolynomialSets identifier) = identifier

getOutputForOneCoefficient :: Integer -> Integer -> Integer -> Integer -> Integer -> Integer
-- returns the output of the polynomial having coefficient coefficient for monomials from degree maxDegree to degree minDegree
--  and monomials of degree < minDegree having total value computedOutput for x value x
--  minDegree must be <= maxDegree
getOutputForOneCoefficient minDegree maxDegree coefficient x computedOutput 
  = List.foldl' (\total degree -> total + (coefficient * x ^ degree)) computedOutput [minDegree .. maxDegree]

getLowerBound :: Integer -> PolynomialConstraints -> Integer -> Integer
-- returns the lower bound for the polynomial having sum totalOutput for monomials of degree <= monomialDegree 
--  and polynomial constraints (PolynomialConstraints x y bound maxPolynomialDegree) 
getLowerBound totalOutput (PolynomialConstraints x y bound maxPolynomialDegree) monomialDegree
  | x >= 0 = getOutputForOneCoefficient (monomialDegree + 1) maxPolynomialDegree (-bound) x totalOutput
  | otherwise = List.foldl' (\total degree ->
                              if degree `mod` 2 == 0
                              then total + (-bound) * x ^ degree
                              else total +   bound  * x ^ degree) totalOutput [monomialDegree + 1 .. maxPolynomialDegree]

checkPolynomialSetPossible :: Integer -> PolynomialConstraints -> Integer -> Bool
-- returns False if it is impossible to produce y for PolynomialConstraints polynomialConstraints, Monomial monomial,
--  and total output computed so far Integer totalOutput
checkPolynomialSetPossible totalOutput (PolynomialConstraints x y bound maxPolynomialDegree) monomialDegree
  | monomialDegree < maxPolynomialDegree =
    let lowerBound = getLowerBound totalOutput (PolynomialConstraints x y bound maxPolynomialDegree) monomialDegree
        upperBound = getOutputForOneCoefficient (monomialDegree + 1) maxPolynomialDegree   bound  x totalOutput in
          lowerBound <= y && y <= upperBound
  | otherwise = totalOutput == y

filterPolynomialsByDegree :: Map.Map Index PolynomialSetWithIdentifier -> Integer -> Map.Map Index PolynomialSetWithIdentifier
-- O(n)
-- returns Map Index PolynomialSetWithIdentifier in which all values have degree correctDegree
filterPolynomialsByDegree mapPolynomialSets correctDegree = Map.filterWithKey (\(Index degree output) polynomialSetID -> degree == correctDegree) mapPolynomialSets

mapValues :: Ord a => Map.Map k a -> Set.Set a
mapValues m = Map.foldlWithKey'(\accumulator key value -> Set.insert value accumulator) Set.empty m

addOneMonomialToPolySets :: Monomial -> PolynomialConstraints -> Map.Map Index PolynomialSetWithIdentifier -> Map.Map Index PolynomialSetWithIdentifier -> Map.Map Index PolynomialSetWithIdentifier

addOneMonomialToPolySets (Monomial coefficient degree) (PolynomialConstraints x y bound polynomialDegree) existingMapDegreeMinus1 existingAccumulatorMap =
  let newMonomial = Monomial coefficient degree -- this does get bound way too early, but it's so small... probably something else?
      newMonomialOutput = coefficient * x ^ degree in
        Map.foldlWithKey' (\accumulator (Index existingDegree existingOutput) !existingPolynomialSetWithID ->
                              let !totalOutput = newMonomialOutput + existingOutput in
                                if checkPolynomialSetPossible totalOutput (PolynomialConstraints x y bound polynomialDegree) degree
                                then
                                  case Map.lookup (Index degree totalOutput) accumulator of
                                    Just (PolynomialSetWithIdentifier !subPolynomialSets _) ->
                                      let !newSubPolynomialSetWithID   = SubPolynomialSetWithIdentifier newMonomial existingPolynomialSetWithID (Identifier $ hash (newMonomial, getPolynomialSetIdentifier existingPolynomialSetWithID)) in
                                        let !newSubPolynomialSets      = List.insert newSubPolynomialSetWithID subPolynomialSets in
                                          let !newPolynomialSetWithID  = PolynomialSetWithIdentifier newSubPolynomialSets (Identifier $ hash $ getSubPolynomialSetsIdentifiers newSubPolynomialSets) in
                                            insert (Index degree totalOutput) newPolynomialSetWithID accumulator
                                    Nothing ->
                                      let !newSubPolynomialSetWithID = SubPolynomialSetWithIdentifier newMonomial existingPolynomialSetWithID (Identifier $ hash (newMonomial, getPolynomialSetIdentifier existingPolynomialSetWithID)) in
                                        let !newPolynomialSetWithID  = PolynomialSetWithIdentifier [newSubPolynomialSetWithID] (Identifier $ hash $ getSubPolynomialSetsIdentifiers [newSubPolynomialSetWithID]) in
                                          insert (Index degree totalOutput) newPolynomialSetWithID accumulator
                                else accumulator
                                ) existingAccumulatorMap existingMapDegreeMinus1

makePolynomialSetsFast :: PolynomialConstraints -> PolynomialsSolutions ()
makePolynomialSetsFast (PolynomialConstraints x y bound degree) = 
  makePolynomialSetsForMaxDegree (PolynomialConstraints x y bound 0) degree

makePolynomialSetsFastDegreeZero :: PolynomialConstraints -> Integer -> PolynomialsSolutions ()
makePolynomialSetsFastDegreeZero (PolynomialConstraints x y bound degree) maxDegree =
  let newMap = (List.foldl' (\accumulator coefficient ->
                let newMonomial         = (Monomial coefficient degree)
                    newSubPolynomialSet = SubPolynomialSetWithIdentifier newMonomial emptyPolynomialSetWithID (Identifier $ hash (newMonomial, getPolynomialSetIdentifier emptyPolynomialSetWithID))
                    newPolynomialSet    = PolynomialSetWithIdentifier [newSubPolynomialSet] (Identifier $ hash $ getSubPolynomialSetsIdentifiers [newSubPolynomialSet]) in
                  Map.insert (Index degree coefficient) newPolynomialSet accumulator)  Map.empty [-bound .. bound]) in
  modify' (const newMap)

makePolynomialSetsForMaxDegree :: PolynomialConstraints -> Integer -> PolynomialsSolutions ()
makePolynomialSetsForMaxDegree (PolynomialConstraints x y bound degree) maxDegree
  | degree == maxDegree =
    if degree == 0
    then makePolynomialSetsFastDegreeZero (PolynomialConstraints x y bound degree) maxDegree
    else do -- base case
      !existingPolysets <- get only need highest level.
      let finalMap = (List.foldl' (\accumulator coefficient ->
                      let newMonomial = (Monomial coefficient degree) in
                        Map.union (addOneMonomialToPolySets newMonomial (PolynomialConstraints x y bound maxDegree) existingPolysets accumulator) accumulator) Map.empty [-bound .. bound])
      ($!!) modify' (const finalMap) -- definitely needs to be deep-seq'ed
  | degree == 0 = do
    makePolynomialSetsFastDegreeZero (PolynomialConstraints x y bound degree) maxDegree
    makePolynomialSetsForMaxDegree (PolynomialConstraints x y bound (degree +1)) maxDegree
  | otherwise = do
    !existingPolysets <- get
    let newMap = (List.foldl' (\accumulator coefficient ->
                  let newMonomial = (Monomial coefficient degree)
                      innerMap = addOneMonomialToPolySets newMonomial (PolynomialConstraints x y bound maxDegree) existingPolysets accumulator in
                    Map.union innerMap accumulator) Map.empty [-bound..bound])
    ($!!) modify' (const newMap)
    makePolynomialSetsForMaxDegree (PolynomialConstraints x y bound (degree +1)) maxDegree

如何解决此内存泄漏问题？

Answer 1

一个问题（也许不是唯一的问题）是($!!) modify' (const finalMap)，在评估时，不深切finalMap。这是因为const finalMap 是一个函数，而 deepseq 不会越过函数。函数的 NFData 实例说：

这个实例是为了方便和与 seq 保持一致。这假设 WHNF 等价于函数的 NF。

而 lambda 抽象已经在 WHNF 中。

例如，这段代码在 ghci 中没有给出任何错误：

Prelude Control.DeepSeq> seq (rnf (const undefined)) ()
()

我会选择put $! rnf finalMap，而不是($!!) modify' (const finalMap)。在评估该表达式时，首先评估 rnf finalMap，这要归功于 $!，然后会触发 deepseq。