Streamline representation and tweak comments

1 files changed, 35 insertions, 43 deletions

diff --git a/src/compiler/GF/Compile/GeneratePMCFG.hs b/src/compiler/GF/Compile/GeneratePMCFG.hs
index d94e67e5d..c3f20a19e 100644
--- a/src/compiler/GF/Compile/GeneratePMCFG.hs
+++ b/src/compiler/GF/Compile/GeneratePMCFG.hs
@@ -600,10 +600,13 @@ data PMCFGEnv = PMCFGEnv !ProdGroups !FunSet
 type ProdGroups = Map.Map (FId,FunId) ProdGroup
 type FunSet     = Map.Map (UArray LIndex SeqId) FunId
 
+-- FIdAlts arrays always use dense 0-based bounds, so their derived Ord matches
+-- the old list ordering of alternatives.
 newtype FIdAlts = FIdAlts (UArray Int FId)
   deriving (Eq, Ord)
 
--- Factors are weighted parameter-choice arrays, in schema traversal order, preserving duplicates.
+-- Factors are weighted parameter-choice arrays, in schema traversal order,
+-- preserving duplicates.
 newtype FIdFactors = FIdFactors [UArray Int FId]
 
 -- Store the exact argument FId alternatives for one Production compactly.
@@ -612,13 +615,13 @@ newtype ArgFIdProduct = ArgFIdProduct [FIdAlts]
 
 -- Accumulator type for Productions with the same result FId and function.
 -- The Set matches the old ProdSet deduplication for argument products.  The
--- optional IntSets cache the per-argument union of FIds when all products have
--- the same arity.  The final Int stores areaSum, the sum of product sizes for
--- distinct products.  A group can be emitted as one compressed Production
--- exactly when the union area equals areaSum.
+-- IntSets cache the per-argument union of FIds.  The final Int stores areaSum,
+-- the sum of product sizes for distinct products.  A group can be emitted as
+-- one compressed Production exactly when the union area equals areaSum.  All
+-- products in a group must have the same argument arity.
 data ProdGroup = ProdGroup
                  !(Set.Set ArgFIdProduct)
-                 !(Maybe [IntSet.IntSet])
+                 ![IntSet.IntSet]
   {-# UNPACK #-} !Int
 
 emptyPMCFGEnv =
@@ -626,31 +629,32 @@ emptyPMCFGEnv =
 
 addFunction :: PMCFGEnv -> FId -> UArray LIndex SeqId -> ArgFIdProduct -> PMCFGEnv
 addFunction (PMCFGEnv prodGroups funSet) !fid fun argProduct =
-  case Map.lookup fun funSet of
-    Just !funid -> PMCFGEnv (insertProduction fid funid argProduct prodGroups) funSet
-    Nothing     -> let !funid = Map.size funSet
-                   in PMCFGEnv (insertProduction fid funid argProduct prodGroups)
-                               (Map.insert fun funid funSet)
+  let (!funid,!funSet') =
+        case Map.lookup fun funSet of
+          Just !funid -> (funid, funSet)
+          Nothing     -> let !funid = Map.size funSet
+                          in (funid, Map.insert fun funid funSet)
+      update Nothing      = Just $! singletonProdGroup argProduct
+      update (Just group) = Just $! insertArgFIdProduct argProduct group
+      !prodGroups' = Map.alter update (fid,funid) prodGroups
+  in PMCFGEnv prodGroups' funSet'
 
 getPMCFG :: PMCFGEnv -> PMCFG
 getPMCFG (PMCFGEnv prodGroups funSet) =
   PMCFG (Map.foldrWithKey addGroup [] prodGroups) (mkSetArray funSet)
   where
     addGroup :: (FId,FunId) -> ProdGroup -> [Production] -> [Production]
-    addGroup (fid,funid) (ProdGroup products mArgSets areaSum) prods
+    addGroup (fid,funid) (ProdGroup products argSets areaSum) prods
       | product (map IntSet.size argSets) == areaSum
                   = Production fid funid (map IntSet.toList argSets) : prods
-      -- We reverse the list for byte-for-byte equivalence with the previous grouping order.
+      -- The old implementation enumerated the global Production Set in ascending
+      -- order, then grouped with Map.fromListWith (++), reversing each group's
+      -- uncompressed products.  Keep that order for stable serialized output.
       | 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
-
 #ifdef PMCFG_TEST_HOOKS
 pmcfgTestBuildPMCFG :: [TestProduction] -> PMCFG
 pmcfgTestBuildPMCFG =
@@ -663,44 +667,31 @@ pmcfgTestBuildPMCFG =
     fIdAltsFromList fids = FIdAlts (listArray (0,length fids-1) fids)
 #endif
 
-insertProduction :: FId -> FunId -> ArgFIdProduct -> ProdGroups -> ProdGroups
-insertProduction !fid !funid argProduct prodGroups =
-  Map.alter update (fid,funid) prodGroups
-  where
-    update Nothing      = Just $! singletonProdGroup argProduct
-    update (Just group) = Just $! insertArgFIdProduct argProduct group
-
 singletonProdGroup :: ArgFIdProduct -> ProdGroup
 singletonProdGroup argProduct@(ArgFIdProduct args) =
   let !products = Set.singleton argProduct
       !argSets  = fmap (insertFIdAlts IntSet.empty) args
       !areaSum  = argFIdProductSize argProduct
-  in ProdGroup products (Just argSets) areaSum
+  in ProdGroup products argSets areaSum
 
 insertArgFIdProduct :: ArgFIdProduct -> ProdGroup -> ProdGroup
-insertArgFIdProduct argProduct group@(ProdGroup products mArgSets areaSum)
+insertArgFIdProduct argProduct group@(ProdGroup products argSets areaSum)
   | Set.member argProduct products
       = group
   | otherwise
       = let !products' = Set.insert argProduct products
-            !mArgSets' = updateArgSets mArgSets argProduct
+            !argSets'  = updateArgSets argSets argProduct
             !areaSum'  = areaSum + argFIdProductSize argProduct
-        in ProdGroup products' mArgSets' areaSum'
+        in ProdGroup products' argSets' areaSum'
   where
     addArgSet argSet fids = insertFIdAlts argSet fids
 
-    updateArgSets Nothing _ = Nothing
-    updateArgSets (Just argSets) (ArgFIdProduct argFIds)
-      | length argSets == length argFIds = let !argSets' = zipWith addArgSet argSets argFIds
-                                           in Just argSets'
-      | otherwise                        = Nothing
-
-argFIdProductArgSets :: Set.Set ArgFIdProduct -> [IntSet.IntSet]
-argFIdProductArgSets products =
-  List.foldl' addProduct (repeat IntSet.empty) (Set.toList products)
-  where
-    addProduct argSets (ArgFIdProduct args) = zipWith addArgSet argSets args
-    addArgSet argSet fids = insertFIdAlts argSet fids
+    updateArgSets argSets (ArgFIdProduct argFIds) = go argSets argFIds
+      where
+        go [] [] = []
+        go (argSet:argSets) (argFIds:argFIds') =
+          addArgSet argSet argFIds : go argSets argFIds'
+        go _ _ = bug "insertArgFIdProduct: arity mismatch"
 
 insertFIdAlts :: IntSet.IntSet -> FIdAlts -> IntSet.IntSet
 insertFIdAlts = foldFIdAlts (\s fid -> IntSet.insert fid s)
@@ -722,12 +713,13 @@ fIdAltsFromFactors factors@(FIdFactors comps)
   where
     !resultSize = fIdFactorsResultSize factors
 
+-- Components are ordered outer-to-inner. This must match the old
+-- reverse (solutions (variants schema) ()) ordering, where the last component
+-- varies fastest.
 fillFIds :: STUArray s Int FId -> Int -> FId -> [UArray Int FId] -> ST s Int
 fillFIds arr !offset !fidAcc [] = do
   writeArray arr offset fidAcc
   return (offset + 1)
--- Components are ordered outer-to-inner. This must match the old
--- reverse (solutions (variants schema) ()) ordering.
 fillFIds arr !offset !fidAcc (factor : factors) =
   foldUArrayM (\offset' fidDelta -> fillFIds arr offset' (fidAcc + fidDelta) factors) offset factor