Merge pull request #5492 from unisonweb/cp/better-branches

unison-runtime/src/Unison/Runtime/MCode.hs

@@ -730,16 +730,33 @@ data GBranch comb
       !(M.Map Text (GSection comb))
   deriving stock (Show, Eq, Ord, Functor, Foldable, Traversable)
 
+branchToEnumMap :: GBranch comb -> Maybe ((GSection comb), EnumMap Word64 (GSection comb))
+branchToEnumMap = \case
+  (Test1 k t d) -> Just (d, EC.mapSingleton k t)
+  (Test2 k1 s1 k2 s2 d) -> Just (d, EC.mapFromList [(k1, s1), (k2, s2)])
+  (TestW d m) -> Just (d, m)
+  _ -> Nothing
+
 -- Convenience patterns for matches used in the algorithms below.
 pattern MatchW :: Int -> (GSection comb) -> EnumMap Word64 (GSection comb) -> (GSection comb)
-pattern MatchW i d cs = Match i (TestW d cs)
+pattern MatchW i d cs <- Match i (branchToEnumMap -> Just (d, cs))
+  where
+    MatchW i d cs = Match i (mkBranch d cs)
 
 pattern MatchT :: Int -> (GSection comb) -> M.Map Text (GSection comb) -> (GSection comb)
 pattern MatchT i d cs = Match i (TestT d cs)
 
 pattern NMatchW ::
   Maybe Reference -> Int -> (GSection comb) -> EnumMap Word64 (GSection comb) -> (GSection comb)
-pattern NMatchW r i d cs = NMatch r i (TestW d cs)
+pattern NMatchW r i d cs <- NMatch r i (branchToEnumMap -> Just (d, cs))
+  where
+    NMatchW r i d cs = NMatch r i $ mkBranch d cs
+
+mkBranch :: (GSection comb) -> (EnumMap Word64 (GSection comb)) -> GBranch comb
+mkBranch d m = case EC.mapToList m of
+  [(k, v)] -> Test1 k v d
+  [(k1, v1), (k2, v2)] -> Test2 k1 v1 k2 v2 d
+  _ -> TestW d m
 
 -- Representation of the variable context available in the current
 -- frame. This tracks tags that have been dumped to the stack for
@@ -1445,7 +1462,7 @@ emitDataMatching ::
   Maybe (ANormal v) ->
   Emit Branch
 emitDataMatching r rns grpr grpn rec ctx cs df =
-  TestW <$> edf <*> traverse (emitCase rns grpr grpn rec ctx) (coerce cs)
+  mkBranch <$> edf <*> traverse (emitCase rns grpr grpn rec ctx) (coerce cs)
   where
     -- Note: this is not really accurate. A default data case needs
     -- stack space corresponding to the actual data that shows up there.
@@ -1489,7 +1506,7 @@ emitRequestMatching rns grpr grpn rec ctx hs df = (,) <$> pur <*> tops
   where
     pur = emitCase rns grpr grpn rec ctx ([BX], df)
     tops = traverse f (coerce hs)
-    f cs = TestW edf <$> traverse (emitCase rns grpr grpn rec ctx) cs
+    f cs = mkBranch edf <$> traverse (emitCase rns grpr grpn rec ctx) cs
     edf = Die "unhandled ability"
 
 emitLitMatching ::

unison-runtime/src/Unison/Runtime/Machine.hs

@@ -2127,7 +2127,7 @@ codeValidate tml cc = do
   ftm <- readTVarIO (freshTm cc)
   rtm0 <- readTVarIO (refTm cc)
   let rs = fst <$> tml
-      rtm = rtm0 `M.withoutKeys` S.fromList rs
+      rtm = rtm0 `M.union` M.fromList (zip rs [ftm ..])
       rns = RN (refLookup "ty" rty) (refLookup "tm" rtm) (const Nothing)
       combinate (n, (r, g)) = evaluate $ emitCombs rns r n g
   (Nothing <$ traverse_ combinate (zip [ftm ..] tml))

unison-runtime/src/Unison/Runtime/Stack.hs

@@ -286,6 +286,11 @@ data GClosure comb
   deriving stock (Show, Functor, Foldable, Traversable)
 {- ORMOLU_ENABLE -}
 
+-- Singleton black hole value to avoid allocation.
+blackHole :: Closure
+blackHole = Closure GBlackHole
+{-# NOINLINE blackHole #-}
+
 pattern PAp :: CombIx -> GCombInfo (RComb Val) -> Seg -> Closure
 pattern PAp cix comb seg = Closure (GPAp cix comb seg)
 
@@ -302,7 +307,9 @@ pattern Captured k a seg = Closure (GCaptured k a seg)
 
 pattern Foreign x = Closure (GForeign x)
 
-pattern BlackHole = Closure GBlackHole
+pattern BlackHole <- Closure GBlackHole
+  where
+    BlackHole = blackHole
 
 pattern UnboxedTypeTag t <- Closure (GUnboxedTypeTag t)
   where