Update GeneratePMCFG.hsoptimize

1 files changed, 16 insertions, 15 deletions

diff --git a/src/compiler/GF/Compile/GeneratePMCFG.hs b/src/compiler/GF/Compile/GeneratePMCFG.hs
index 5a0359a91..22188b26f 100644
--- a/src/compiler/GF/Compile/GeneratePMCFG.hs
+++ b/src/compiler/GF/Compile/GeneratePMCFG.hs
@@ -599,6 +599,7 @@ type FunSet     = Map.Map (UArray LIndex SeqId) FunId
 newtype FIdAlts = FIdAlts (UArray Int FId)
   deriving (Eq, Ord)
 
+-- Factors are weighted parameter-choice arrays, in schema traversal order, preserving duplicates.
 newtype FIdFactors = FIdFactors [UArray Int FId]
 
 -- Keep exact argument FId products to preserve the old finalizer's duplicate
@@ -639,6 +640,10 @@ getPMCFG (PMCFGEnv prodGroups funSet) =
       -- We reverse the list for byte-to-byte equivalence with the previous grouping order.
       | otherwise = map (Production fid funid . unpackArgFIdProduct) (reverse (Set.toList products)) ++ prods
       where
+        unpackArgFIdProduct :: ArgFIdProduct -> [[FId]]
+        unpackArgFIdProduct (ArgFIdProduct args) = map fidAltsToList args
+
+        argSets :: [IntSet.IntSet]
         argSets = case mArgSets of
                     Just argSets -> argSets
                     Nothing      -> argFIdProductArgSets products
@@ -665,9 +670,9 @@ insertProduction !fid !funid argProduct prodGroups =
         Just group -> insertArgFIdProduct argProduct group
 
 singletonProdGroup :: ArgFIdProduct -> ProdGroup
-singletonProdGroup argProduct =
+singletonProdGroup argProduct@(ArgFIdProduct args) =
   let !products = Set.singleton argProduct
-      !argSets  = argFIdProductArgSetsOne argProduct
+      !argSets  = fmap (insertFIdAlts IntSet.empty) args
       !areaSum  = argFIdProductSize argProduct
   in ProdGroup products (Just argSets) areaSum
 
@@ -681,7 +686,7 @@ insertArgFIdProduct argProduct group@(ProdGroup products mArgSets areaSum)
             !areaSum'  = areaSum + argFIdProductSize argProduct
         in ProdGroup products' mArgSets' areaSum'
   where
-    addArgSet argSet fids = foldFIdAlts (\s fid -> IntSet.insert fid s) argSet fids
+    addArgSet argSet fids = insertFIdAlts argSet fids
 
     updateArgSets Nothing _ = Nothing
     updateArgSets (Just argSets) (ArgFIdProduct argFIds)
@@ -694,14 +699,10 @@ argFIdProductArgSets products =
   List.foldl' addProduct (repeat IntSet.empty) (reverse (Set.toList products))
   where
     addProduct argSets (ArgFIdProduct args) = zipWith addArgSet argSets args
-    addArgSet argSet fids = foldFIdAlts (\s fid -> IntSet.insert fid s) argSet fids
-
-argFIdProductArgSetsOne :: ArgFIdProduct -> [IntSet.IntSet]
-argFIdProductArgSetsOne (ArgFIdProduct args) =
-  fmap (foldFIdAlts (\s fid -> IntSet.insert fid s) IntSet.empty) args
+    addArgSet argSet fids = insertFIdAlts argSet fids
 
-unpackArgFIdProduct :: ArgFIdProduct -> [[FId]]
-unpackArgFIdProduct (ArgFIdProduct args) = map fidAltsToList args
+insertFIdAlts :: IntSet.IntSet -> FIdAlts -> IntSet.IntSet
+insertFIdAlts = foldFIdAlts (\s fid -> IntSet.insert fid s)
 
 argFIdProductSize :: ArgFIdProduct -> Int
 argFIdProductSize (ArgFIdProduct args) = product (map fidAltsSize args)
@@ -721,13 +722,13 @@ fIdAltsFromFactors factors@(FIdFactors comps)
     !resultSize = fIdFactorsResultSize factors
 
 fillFIds :: STUArray s Int FId -> Int -> FId -> [UArray Int FId] -> ST s Int
-fillFIds arr !offset !sum [] = do
-  writeArray arr offset sum
-  return (offset+1)
+fillFIds arr !offset !acc [] = do
+  writeArray arr offset acc
+  return (offset + 1)
 -- Components are ordered outer-to-inner. This must match the old
 -- reverse (solutions (variants schema) ()) ordering.
-fillFIds arr !offset !sum (choices:choices') =
-  foldUArrayM (\offset' choice -> fillFIds arr offset' (sum+choice) choices') offset choices
+fillFIds arr !offset !acc (choices : choices') =
+  foldUArrayM (\offset' choice -> fillFIds arr offset' (acc + choice) choices') offset choices
 
 foldUArrayM :: Monad m => (a -> FId -> m a) -> a -> UArray Int FId -> m a
 foldUArrayM f z arr = go (fst bnds) z