Fix combination bug

src/Dodge/Combine.hs

@@ -2,11 +2,13 @@
 module Dodge.Combine where
 import Dodge.Base.You
 import Dodge.Data
+import Dodge.Item.Amount
 --import Dodge.Combine.Data
 import Multiset
 import Dodge.Item.Weapon.BulletGuns
 import Dodge.Item.Weapon.Launcher
 --import Dodge.Item.Craftable
+import SimpleTrie
 
 import Control.Lens
 import Data.Bifunctor
@@ -14,10 +16,17 @@ import qualified Data.IntMap.Strict as IM
 --import qualified Data.IntSet as IS
 import qualified Data.Map.Strict as M
 import Data.Maybe
-import Data.List (scanl')
+import Data.List (sort,scanl')
 
 itemCombinations :: [ (M.Map CombineType Int, Item) ]
-itemCombinations = map (first toMultiset) $
+itemCombinations = map (first toMultiset) itemCombinations'
+
+combinationsTrie :: Trie CombineType Item
+combinationsTrie = foldr (uncurry insertInTrie) emptyTrie
+    $ map (first sort) itemCombinations'
+
+itemCombinations' :: [([CombineType],Item)]
+itemCombinations' =
     [ p [PIPE,HARDWARE] (bangStick 1)
     , p [BANGSTICK 1,TIN] pistol
     , p [PISTOL, SPRING, HARDWARE] autoPistol
@@ -37,6 +46,7 @@ itemCombinations = map (first toMultiset) $
     , p [BANGCANE,TIN] rifle
     , p [RIFLE,SPRING,HARDWARE] autoRifle
     , p [RIFLE,SPRING,CAN] assaultRifle
+    , p [BANGCANE,BANGCANE] miniGun
 
     , p [PIPE,PIPE,PIPE,HARDWARE] bangRod
     , p [BANGROD,PLANK,HARDWARE]  elephantGun
@@ -52,28 +62,65 @@ itemCombinations = map (first toMultiset) $
   where
     p = (,)
 
-combinations :: [(M.Map CombineType Int, CombineType)]
-combinations = map (second _itType) itemCombinations
-
 combineToItem :: M.Map CombineType Int -> Maybe Item
 combineToItem m = lookup m itemCombinations
 
+expandItem :: (Int,Item) -> [(Int,CombineType)]
+expandItem (i,it) = case it ^? itConsumption . itAmount of
+    Nothing -> [(i,_itType it)]
+    Just j  -> replicate j (i,_itType it)
+
 combineItemList :: IM.IntMap Item -> [([Int],Item)]
-combineItemList m = mapMaybe (pushoutmaybe . second combineToItem)
+combineItemList m 
+    = mapMaybe (pushoutmaybe . second combineToItem)
     . concatMap itemsMultisets
     . powlistN 4
     $ IM.toList m
   where
     pushoutmaybe (_,Nothing) = Nothing
     pushoutmaybe (x,Just y)  = Just (x,y)
---combineItemList :: IM.IntMap CombineType -> [([Int],Item)]
---combineItemList m = mapMaybe (pushoutmaybe . second (combineToItem . toMultiset) . unzip)
---    . powlistN 4
---    $ IM.toList m
---  where
---    pushoutmaybe (_,Nothing) = Nothing
---    pushoutmaybe (x,Just y)  = Just (x,y)
 
+lookupTypeTrie :: [CombineType] -> Maybe Item
+lookupTypeTrie = flip lookupTrie combinationsTrie . sort
+
+combineItemList' :: IM.IntMap Item -> [([Int],Item)]
+combineItemList'
+    = mapMaybe (pushoutmaybe . second lookupTypeTrie)
+    . listCombos
+  where
+    pushoutmaybe (_,Nothing) = Nothing
+    pushoutmaybe (x,Just y)  = Just (x,y)
+
+listCombos :: IM.IntMap Item -> [([Int], [CombineType])]
+listCombos = concatMap expandOut . invertListInvMult
+
+expandOut :: [(CombineType,Int,Int)] -> [([Int],[CombineType])]
+expandOut [] = [([],[])]
+expandOut (x:xs) = (<>) <$> f x <*> expandOut xs
+  where
+    f (ct,i,n) = map (\j -> (replicate j i,replicate j ct)) [1..n]
+
+invertListInvMult :: IM.IntMap Item -> [[(CombineType,Int,Int)]]
+invertListInvMult = powlistN 4 . invertListInv
+
+invertListInv :: IM.IntMap Item -> [(CombineType,Int,Int)]
+invertListInv = IM.foldrWithKey
+    (\k it -> ((_itType it, k, itStackAmount it) :) )
+    []
+
+invertInv :: IM.IntMap Item -> M.Map CombineType (IM.IntMap Int)
+invertInv = IM.foldrWithKey
+    (\k it -> M.insertWith IM.union (_itType it) (IM.singleton k (itStackAmount it)))
+    M.empty
+
+expandCounts :: (Int,(Int,a)) -> [[(Int,a)]]
+expandCounts (i,(j,ct)) = [ replicate k (i,ct) | k <- [1..j]]
+
+multisetsWithCounts :: IM.IntMap Item -> [[(Int,(Int,CombineType))]]
+multisetsWithCounts = powlistN 4 . IM.toList . IM.map countAndType
+  where
+    countAndType it = (fromMaybe 1 $ it ^? itConsumption . itAmount
+        , _itType it)
 
 -- multiple multisets produced by a single item stack are combined using the list monad
 itemsMultisets :: [(Int,Item)] -> [([Int],M.Map CombineType Int)]
@@ -88,18 +135,8 @@ itemMultisets (i,it) = case it ^? itConsumption . itAmount of
   where
     thetype = _itType it
 
-itemTypeAmounts :: Item -> [[CombineType]]
-itemTypeAmounts it = case it ^? itConsumption . itAmount of
-    Nothing -> [[thetype]]
-    Just i  -> map (`replicate` thetype) [1 .. i]
-  where
-    thetype = _itType it
-
 combineItemListYou :: World -> [([Int],Item)]
-combineItemListYou = combineItemList . yourInv
-
-combineListYou :: World -> [([Int],CombineType)]
-combineListYou = fmap (second _itType) . combineItemListYou
+combineItemListYou = combineItemList' . yourInv
 
 toggleCombineInv :: World -> World
 toggleCombineInv w = case _inventoryMode w of
@@ -109,10 +146,7 @@ toggleCombineInv w = case _inventoryMode w of
 enterCombineInv :: World -> World
 enterCombineInv w = w & inventoryMode .~ CombineInventory mi
   where
-    mi = 0 <$ listToMaybe (combineListYou w)
-
-combine :: M.Map CombineType Int -> Maybe CombineType
-combine m = lookup m combinations
+    mi = 0 <$ listToMaybe (combineItemListYou w)
 
 combineSizes :: World -> [Int]
 combineSizes = map (ceiling . _itInvSize . snd) . combineItemListYou