Allow comptime expressions as match labels

src/Solcore/Backend/EmitHull.hs

@@ -371,6 +371,7 @@ emitWordMatch scrutinee alts = do
       hullStmts <- emitStmts stmts
       let hullName = show n
       return (Hull.Alt (Hull.PVar hullName) "$_" hullStmts)
+    emitWordAlt _ (MastPExp _, _) = errorsEM ["PANIC: MastPExp reached EmitHull — was not evaluated by MastEval"]
     emitWordAlt _ (pat, _) = errorsEM ["emitWordAlt not implemented for", show pat]
 
 type BranchMap = Map.Map Name [Hull.Stmt]

src/Solcore/Backend/Mast.hs

@@ -125,6 +125,7 @@ data MastPat
   | MastPCon MastId [MastPat]
   | MastPWildcard
   | MastPLit Literal
+  | MastPExp MastExp -- comptime expression label; must be evaluated by MastEval
   deriving (Eq, Ord, Show)
 
 -----------------------------------------------------------------------
@@ -269,6 +270,7 @@ instance Pretty MastPat where
     | otherwise = ppr n >< parens (commaSep $ map ppr ps)
   ppr MastPWildcard = text "_"
   ppr (MastPLit l) = ppr l
+  ppr (MastPExp e) = text "comptime" <+> ppr e
 
 -----------------------------------------------------------------------
 -- Helpers (shared with SolcorePretty)

src/Solcore/Backend/MastEval.hs

@@ -196,8 +196,19 @@ evalAlt :: VEnv -> MastAlt -> EvalM MastAlt
 evalAlt env (pat, body) = do
   -- Pattern bindings shadow existing bindings, but we don't track them
   -- (conservative: treat all pattern-bound vars as unknown)
+  pat' <- evalPat env pat
   (_, body') <- evalStmts env body
-  pure (pat, body')
+  pure (pat', body')
+
+-- Evaluate expression labels in patterns.
+-- MastPExp must reduce to a literal; any other form is a compile-time error.
+evalPat :: VEnv -> MastPat -> EvalM MastPat
+evalPat env (MastPExp e) = do
+  e' <- evalExp env e
+  case e' of
+    MastLit l -> pure (MastPLit l)
+    _ -> error $ "comptime expression in match label could not be evaluated to a literal: " ++ show e'
+evalPat _ pat = pure pat
 
 -----------------------------------------------------------------------
 -- Evaluate expressions
@@ -333,7 +344,8 @@ evalFunBody env (stmt : rest) = case stmt of
     pure $ if isKnownValue e' then Just e' else Nothing
   MastMatch scrut alts -> do
     scrut' <- evalExp env scrut
-    case matchAlts env scrut' alts of
+    alts' <- mapM (\(p, b) -> (,b) <$> evalPat env p) alts
+    case matchAlts env scrut' alts' of
       Just (env', body) -> evalFunBody env' body
       Nothing -> pure Nothing -- Scrutinee not known, can't select branch
   MastAsm _ -> pure Nothing -- Should not happen: purity analysis excludes asm functions
@@ -373,6 +385,7 @@ findLitMatch env lit ((pat, body) : rest) =
       let env' = Map.insert varId (MastLit lit) env
        in Just (env', body)
     MastPWildcard -> Just (env, body)
+    MastPExp _ -> error "PANIC: MastPExp reached findLitMatch — evalAlt failed to evaluate it"
     _ -> findLitMatch env lit rest
 
 -- Bind pattern variables to argument expressions
@@ -401,6 +414,7 @@ bindPatterns env (pat : pats) (arg : args) =
       case arg of
         MastLit argLit | argLit == patLit -> bindPatterns env pats args
         _ -> Nothing
+    MastPExp _ -> error "PANIC: MastPExp reached bindPatterns — evalAlt failed to evaluate it"
 
 -----------------------------------------------------------------------
 -- Helpers

src/Solcore/Backend/Specialise.hs

@@ -499,8 +499,15 @@ specMatch exps alts = do
       -- debug ["specAlt, pattern: ", show pat]
       -- debug ["specAlt, body: ", show body]
       body' <- specBody body
-      pat' <- atCurrentSubst pat
+      pat' <- mapM specPat =<< atCurrentSubst pat
       return (pat', body')
+    -- Specialize function calls inside PExp patterns.
+    -- Other pattern forms only need type substitution (handled by atCurrentSubst).
+    specPat :: Pat Id -> SM (Pat Id)
+    specPat (PExp e) = do
+      ty <- atCurrentSubst (typeOfTcExp e)
+      PExp <$> specExp e ty
+    specPat p = pure p
     specScruts = mapM specScrut
     specScrut e = do
       ty <- atCurrentSubst (typeOfTcExp e)
@@ -853,3 +860,4 @@ toMastPat (PVar i) = MastPVar (toMastId i)
 toMastPat (PCon i ps) = MastPCon (toMastId i) (map toMastPat ps)
 toMastPat PWildcard = MastPWildcard
 toMastPat (PLit l) = MastPLit l
+toMastPat (PExp e) = MastPExp (toMastExp e)

src/Solcore/Desugarer/DecisionTreeCompiler.hs

@@ -141,12 +141,12 @@ compileMatrix tys occs matrix bacts = do
     (testTy : restTys, testOcc : restOccs, _) -> do
       let firstCol = [p | (p : _) <- matrix']
           headCons = nub (mapMaybe patHeadCon firstCol)
-          headLits = nub (mapMaybe patHeadLit firstCol)
-      case (headCons, headLits) of
+          headAtomics = nub (mapMaybe patHeadAtomic firstCol)
+      case (headCons, headAtomics) of
         (c : cs, _) ->
           buildConSwitch testOcc restOccs testTy restTys matrix' bacts (c : cs)
-        ([], ls@(_ : _)) ->
-          buildLitSwitch testOcc restOccs testTy restTys matrix' bacts ls
+        ([], as@(_ : _)) ->
+          buildAtomicSwitch testOcc restOccs testTy restTys matrix' bacts as
         ([], []) ->
           -- All wildcards after reordering: strip column and continue
           compileMatrix restTys restOccs (defaultMatrix matrix') (defaultBoundActs testOcc matrix' bacts)
@@ -220,34 +220,34 @@ buildConSwitch testOcc restOccs _testTy restTys matrix bacts headCons = do
               pure (Just Fail)
             else Just <$> compileMatrix restTys restOccs defMat defBacts
 
-buildLitSwitch ::
+buildAtomicSwitch ::
   Occurrence ->
   [Occurrence] ->
   Ty ->
   [Ty] ->
   PatternMatrix ->
   [BoundAction] ->
-  [Literal] ->
+  [AtomicPat] ->
   CompilerM DecisionTree
-buildLitSwitch testOcc restOccs testTy restTys matrix bacts headLits = do
-  branches <- mapM buildBranch headLits
+buildAtomicSwitch testOcc restOccs testTy restTys matrix bacts headAtomics = do
+  branches <- mapM buildBranch headAtomics
   let defMat = defaultMatrix matrix
       defBacts = defaultBoundActs testOcc matrix bacts
   mDefault <-
     if null defMat
       then do
-        litWit <- freshPat testTy
+        wit <- freshPat testTy
         restWit <- mapM freshPat restTys
         ctx <- askWarnCtx
-        tell [NonExhaustive ctx (litWit : restWit)]
+        tell [NonExhaustive ctx (wit : restWit)]
         pure (Just Fail)
       else Just <$> compileMatrix restTys restOccs defMat defBacts
-  pure (LitSwitch testOcc branches mDefault)
+  pure (AtomicSwitch testOcc branches mDefault)
   where
-    buildBranch lit = do
-      let specMat = specializeLit lit matrix
-          specBacts = litSpecializedBoundActs lit testOcc matrix bacts
-      (lit,) <$> compileMatrix restTys restOccs specMat specBacts
+    buildBranch apat = do
+      let specMat = specializeAtomic apat matrix
+          specBacts = atomicSpecializedBoundActs apat testOcc matrix bacts
+      (apat,) <$> compileMatrix restTys restOccs specMat specBacts
 
 instance Compile (Exp Id) where
   compile v@(Var _) =
@@ -303,9 +303,9 @@ treeToStmt occMap (Switch occ branches mDef) = do
       v <- freshPat ty
       pure [([v], body)]
   pure (Match [scrutinee] (eqns ++ defEqns))
-treeToStmt occMap (LitSwitch occ branches mDef) = do
+treeToStmt occMap (AtomicSwitch occ branches mDef) = do
   scrutinee <- occToExp occMap occ
-  eqns <- mapM (litBranchToEqn occMap) branches
+  eqns <- mapM (atomicBranchToEqn occMap) branches
   defEqns <- case mDef of
     Nothing -> pure []
     Just def -> do
@@ -349,10 +349,13 @@ conBranchToEqn occMap occ (k, srcPats, tree) = do
   body <- treeToBody occMap' tree
   pure ([PCon k srcPats], body)
 
-litBranchToEqn :: OccMap -> (Literal, DecisionTree) -> CompilerM (PatternRow, Action)
-litBranchToEqn occMap (lit, tree) = do
+atomicBranchToEqn :: OccMap -> (AtomicPat, DecisionTree) -> CompilerM (PatternRow, Action)
+atomicBranchToEqn occMap (Left lit, tree) = do
   body <- treeToBody occMap tree
   pure ([PLit lit], body)
+atomicBranchToEqn occMap (Right e, tree) = do
+  body <- treeToBody occMap tree
+  pure ([PExp e], body)
 
 -- definition of a monad
 
@@ -445,9 +448,14 @@ data DecisionTree
   = Leaf [(Id, Occurrence)] Action -- var-occ bindings to substitute before running action
   | Fail
   | Switch Occurrence [(Id, [Pattern], DecisionTree)] (Maybe DecisionTree)
-  | LitSwitch Occurrence [(Literal, DecisionTree)] (Maybe DecisionTree)
+  | AtomicSwitch Occurrence [(AtomicPat, DecisionTree)] (Maybe DecisionTree)
   deriving (Eq, Show)
 
+-- An atomic pattern is a compile-time value used as a match label.
+-- Left  = integer literal (already evaluated)
+-- Right = comptime expression (evaluated post-specialization by MastEval)
+type AtomicPat = Either Literal (Exp Id)
+
 -- data constructor information and environment
 
 type TypeEnv = Map Name ConInfo
@@ -529,20 +537,22 @@ specRow k pats (p : rest) =
       | otherwise -> Nothing
     PVar _ -> Just (pats ++ rest)
     PLit _ -> Nothing
+    PExp _ -> Nothing
     _ -> error "PANIC! Found wildcard in specRow"
 
-specializeLit :: Literal -> PatternMatrix -> PatternMatrix
-specializeLit lit = mapMaybe specLitRow
+specializeAtomic :: AtomicPat -> PatternMatrix -> PatternMatrix
+specializeAtomic apat = mapMaybe specAtomicRow
   where
-    specLitRow [] = Nothing
-    specLitRow (p : rest) =
+    specAtomicRow [] = Nothing
+    specAtomicRow (p : rest) =
       case p of
-        PLit l
-          | l == lit -> Just rest
-          | otherwise -> Nothing
+        PLit l | Left l == apat -> Just rest
+        PLit _ -> Nothing
+        PExp e | Right e == apat -> Just rest
+        PExp _ -> Nothing
         PVar _ -> Just rest
         PCon _ _ -> Nothing
-        _ -> error "PANIC! Found wildcard in specializeLit"
+        _ -> error "PANIC! Found wildcard in specializeAtomic"
 
 -- matrix definitions
 
@@ -555,6 +565,7 @@ defaultMatrix = concatMap defaultRow
         PVar _ -> [rest]
         PCon _ _ -> []
         PLit _ -> []
+        PExp _ -> []
         _ -> error "PANIC! Found wildcard in defaultMatrix"
 
 specializedBoundActs :: Id -> Occurrence -> PatternMatrix -> [BoundAction] -> [BoundAction]
@@ -569,6 +580,7 @@ specializedBoundActs k testOcc rows bacts =
       PCon k' _ -> idName k' == idName k
       PVar _ -> True
       PLit _ -> False
+      PExp _ -> False
       PWildcard -> error "PANIC! Found wildcard in specializedBoundActs"
     addVarBinding (PVar v : _) binds = binds ++ [(v, testOcc)]
     addVarBinding _ binds = binds
@@ -585,16 +597,17 @@ defaultBoundActs testOcc rows bacts =
     addVarBinding (PVar v : _) binds = binds ++ [(v, testOcc)]
     addVarBinding _ binds = binds
 
-litSpecializedBoundActs :: Literal -> Occurrence -> PatternMatrix -> [BoundAction] -> [BoundAction]
-litSpecializedBoundActs lit testOcc rows bacts =
+atomicSpecializedBoundActs :: AtomicPat -> Occurrence -> PatternMatrix -> [BoundAction] -> [BoundAction]
+atomicSpecializedBoundActs apat testOcc rows bacts =
   [ (addVarBinding row binds, a)
     | (row, (binds, a)) <- zip rows bacts,
-      rowMatchesLit row
+      rowMatchesAtomic row
   ]
   where
-    rowMatchesLit [] = False
-    rowMatchesLit (p : _) = case p of
-      PLit l -> l == lit
+    rowMatchesAtomic [] = False
+    rowMatchesAtomic (p : _) = case p of
+      PLit l -> Left l == apat
+      PExp e -> Right e == apat
       PVar _ -> True
       _ -> False
     addVarBinding (PVar v : _) binds = binds ++ [(v, testOcc)]
@@ -662,9 +675,10 @@ patHeadCon :: Pattern -> Maybe Id
 patHeadCon (PCon k _) = Just k
 patHeadCon _ = Nothing
 
-patHeadLit :: Pattern -> Maybe Literal
-patHeadLit (PLit l) = Just l
-patHeadLit _ = Nothing
+patHeadAtomic :: Pattern -> Maybe AtomicPat
+patHeadAtomic (PLit l) = Just (Left l)
+patHeadAtomic (PExp e) = Just (Right e)
+patHeadAtomic _ = Nothing
 
 -- redundancy checking (pre-pass over the original matrix)
 
@@ -692,7 +706,9 @@ isUseful tys matrix (q1 : qRest) =
       specMat <- specialize k fieldTys matrix
       isUseful (fieldTys ++ drop 1 tys) specMat (ps ++ qRest)
     PLit l ->
-      isUseful (drop 1 tys) (specializeLit l matrix) qRest
+      isUseful (drop 1 tys) (specializeAtomic (Left l) matrix) qRest
+    PExp _ ->
+      pure True -- conservative: expression labels are always considered useful
     PVar _ ->
       case tys of
         [] -> pure False
@@ -718,10 +734,10 @@ inhabitsM _ [] = pure Nothing
 inhabitsM matrix (ty : restTys) = do
   let firstCol = [p | (p : _) <- matrix]
       headCons = nub (mapMaybe patHeadCon firstCol)
-      headLits = nub (mapMaybe patHeadLit firstCol)
-  case (headCons, headLits) of
+      headAtomics = nub (mapMaybe patHeadAtomic firstCol)
+  case (headCons, headAtomics) of
     (cs@(_ : _), _) -> inhabitsConCol matrix ty restTys cs
-    ([], _ : _) -> inhabitsLitCol matrix ty restTys
+    ([], _ : _) -> inhabitsAtomCol matrix ty restTys
     ([], []) -> do
       r <- inhabitsM (defaultMatrix matrix) restTys
       case r of
@@ -764,8 +780,8 @@ inhabitsConCol matrix _ty restTys headCons =
           let (fieldWit, restWit) = splitAt (length fieldTys) wit
           pure (Just (PCon k fieldWit : restWit))
 
-inhabitsLitCol :: PatternMatrix -> Ty -> [Ty] -> CompilerM (Maybe [Pattern])
-inhabitsLitCol matrix ty restTys = do
+inhabitsAtomCol :: PatternMatrix -> Ty -> [Ty] -> CompilerM (Maybe [Pattern])
+inhabitsAtomCol matrix ty restTys = do
   r <- inhabitsM (defaultMatrix matrix) restTys
   case r of
     Nothing -> pure Nothing

src/Solcore/Desugarer/IfDesugarer.hs

@@ -57,6 +57,7 @@ desugarBoolPat (PCon c@(Id n _) ps)
 desugarBoolPat (PVar a) = PVar a
 desugarBoolPat PWildcard = PWildcard
 desugarBoolPat (PLit l) = PLit l
+desugarBoolPat (PExp e) = PExp e
 
 -- desugaring the boolean type constructor
 

src/Solcore/Frontend/Parser/SolcoreParser.y

@@ -384,6 +384,7 @@ Pattern :: { Pat }
 Pattern : Name PatternList                         {Pat $1 $2}
         | '_'                                      {PWildcard}
         | Literal                                  {PLit $1}
+        | 'comptime' Expr                          {PExp $2}
         | '(' Pattern ')'                          {$2}
         | PatternList                              {Pat (Name "pair") $1}
 

src/Solcore/Frontend/Pretty/SolcorePretty.hs

@@ -318,6 +318,8 @@ instance (Pretty a) => Pretty (Pat a) where
     text "_"
   ppr (PLit l) =
     ppr l
+  ppr (PExp e) =
+    text "comptime" <+> ppr e
 
 instance Pretty Literal where
   ppr (IntLit l) = integer (toInteger l)

src/Solcore/Frontend/Pretty/TreePretty.hs

@@ -330,6 +330,8 @@ instance Pretty Pat where
     text "_"
   ppr (PLit l) =
     ppr l
+  ppr (PExp e) =
+    text "comptime" <+> ppr e
 
 instance Pretty Literal where
   ppr (IntLit l) = integer (toInteger l)

src/Solcore/Frontend/Syntax/NameResolution.hs

@@ -255,6 +255,7 @@ instance Resolve S.Pat where
 
   resolve S.PWildcard = pure PWildcard
   resolve (S.PLit l) = PLit <$> resolve l
+  resolve (S.PExp e) = PExp <$> resolve e
   resolve p@(S.Pat n ps) =
     do
       ps' <- resolve ps `wrapError` p

src/Solcore/Frontend/Syntax/Stmt.hs

@@ -52,6 +52,7 @@ data Pat a
   | PCon a [Pat a]
   | PWildcard
   | PLit Literal
+  | PExp (Exp a) -- comptime expression label (numeric matches only)
   deriving (Eq, Ord, Show, Data, Typeable)
 
 -- definition of literals

src/Solcore/Frontend/Syntax/SyntaxTree.hs

@@ -240,6 +240,7 @@ data Pat
   = Pat Name [Pat]
   | PWildcard
   | PLit Literal
+  | PExp Exp -- comptime expression label (numeric matches only)
   deriving (Eq, Ord, Show, Data, Typeable)
 
 -- definition of literals

src/Solcore/Frontend/TypeInference/Erase.hs

@@ -98,3 +98,5 @@ instance Erase (Pat Id) where
     PWildcard
   erase (PLit l) =
     PLit l
+  erase (PExp e) =
+    PExp (erase e)