3 files changed, 45 insertions, 22 deletions

diff --git a/subprojects/frontend/src/language/folding.ts b/subprojects/frontend/src/language/folding.ts
index b4d4ca22..0afabfa0 100644
--- a/subprojects/frontend/src/language/folding.ts
+++ b/subprojects/frontend/src/language/folding.ts
@@ -109,7 +109,7 @@ function foldWithSibling(node: SyntaxNode): FoldRange | null {
   return null;
 }
 
-function foldWholeNode(node: SyntaxNode): FoldRange {
+export function foldWholeNode(node: SyntaxNode): FoldRange {
   return {
     from: node.from,
     to: node.to,
diff --git a/subprojects/frontend/src/language/problem.grammar b/subprojects/frontend/src/language/problem.grammar
index 7949f90c..0d3f0e3a 100644
--- a/subprojects/frontend/src/language/problem.grammar
+++ b/subprojects/frontend/src/language/problem.grammar
@@ -63,11 +63,11 @@ statement {
   //   kw<"fn"> RelationName RelationName ParameterList<Parameter>?
   //   FunctionBody { ("=" sep<OrOp, Case>)? "." }
   // } |
-  //RuleDefinition {
-  //  kw<"rule">
-  //  RuleName ParameterList<Parameter>?
-  //  RuleBody { ":" sep<OrOp, Conjunction> "==>" sep<OrOp, Consequent> "." }
-  //} |
+  RuleDefinition {
+    (ckw<"decision"> | ckw<"propagation">)? kw<"rule">
+    RuleName ParameterList<Parameter>?
+    RuleBody { ":" sep<OrOp, Conjunction> "==>" sep<OrOp, Consequent> "." }
+  } |
   AtomDeclaration {
     kw<"declare">? ckw<"atom"> sep<",", AtomNodeName> "."
   } |
@@ -96,7 +96,12 @@ FeatureDeclaration {
   ";"?
 }
 
-Parameter { RelationName? VariableName }
+Parameter {
+  kw<"candidate">? (kw<"may"> | kw<"must">)?
+  RelationName?
+  (AndBinding | StarBinding)?
+  VariableName
+}
 
 // Use @dynamicPrecedence to prevent a(b) from being parsed as Expr { a } Expr { b }
 // instead of Atom { a(b) }
@@ -125,7 +130,12 @@ BinaryExpr {
 }
 
 UnaryExpr {
-  !prefix ("+" | "-" | "!" | kw<"count">) Expr
+  !prefix
+  (
+    "+" | "-" | "!" | kw<"count"> |
+    kw<"candidate"> | kw<"may"> | kw<"must">
+  )
+  Expr
 }
 
 CastExpr { !cast Expr kw<"as"> DatatypeName }
@@ -136,16 +146,18 @@ Aggregation {
 
 Atom { RelationName "+"? ParameterList<Expr> }
 
-//Consequent { ("," | Action)+ }
+Consequent { ("," | Action)+ }
 
-//Action {
-//   ckw<"new"> VariableName ("<:" VariableName)? |
-//   kw<"delete"> VariableName |
-//   Literal
-//}
+Action {
+  (NotOp | UnknownOp)? RelationName
+  ParameterList<AssertionActionArgument>
+  (":" Expr)?
+}
 
 AssertionArgument { NodeName | StarArgument }
 
+AssertionActionArgument { VariableName | StarArgument }
+
 Constant { Real | String | StarMult | LogicValue }
 
 ReferenceKind {
@@ -170,7 +182,7 @@ RelationName { QualifiedName ~name }
 
 DatatypeName { QualifiedName }
 
-//RuleName { QualifiedName }
+RuleName { QualifiedName }
 
 AtomNodeName { QualifiedName }
 
@@ -247,5 +259,9 @@ sep1<separator, content> { content (separator content)* }
 
   StarArgument { "*" }
 
+  StarBinding { "*" }
+
+  AndBinding { "&" }
+
   "{" "}" "(" ")" "[" "]" "." ".." "," ":" "->" "<->" "+" "-" "**" "=" "+="
 }
diff --git a/subprojects/frontend/src/language/problemLanguageSupport.ts b/subprojects/frontend/src/language/problemLanguageSupport.ts
index dd5d6347..ae998d20 100644
--- a/subprojects/frontend/src/language/problemLanguageSupport.ts
+++ b/subprojects/frontend/src/language/problemLanguageSupport.ts
@@ -14,7 +14,12 @@ import {
 } from '@codemirror/language';
 import { styleTags, tags as t } from '@lezer/highlight';
 
-import { foldBlockComment, foldConjunction, foldDeclaration } from './folding';
+import {
+  foldBlockComment,
+  foldConjunction,
+  foldDeclaration,
+  foldWholeNode,
+} from './folding';
 import {
   indentBlockComment,
   indentDeclaration,
@@ -30,11 +35,13 @@ const parserWithMetadata = parser.configure({
       'module problem class enum pred fn scope': t.definitionKeyword,
       'import as declare atom multi': t.definitionKeyword,
       'extern datatype aggregator': t.definitionKeyword,
+      rule: t.definitionKeyword,
       'abstract extends refers contains container opposite': t.modifier,
       default: t.modifier,
+      'propagation decision': t.modifier,
       'true false unknown error': t.keyword,
+      'candidate may must': t.operatorKeyword,
       'count in is': t.operatorKeyword,
-      // 'new delete': t.keyword,
       NotOp: t.operator,
       UnknownOp: t.operator,
       OrOp: t.separator,
@@ -45,7 +52,7 @@ const parserWithMetadata = parser.configure({
       'RelationName/QualifiedName': t.typeName,
       'DatatypeName/QualifiedName': t.keyword,
       'AggregatorName/QualifiedName': t.operatorKeyword,
-      // 'RuleName/QualifiedName': t.typeName,
+      'RuleName/QualifiedName': t.typeName,
       'AtomNodeName/QualifiedName': t.atom,
       'VariableName/QualifiedName': t.variableName,
       'ModuleName/QualifiedName': t.typeName,
@@ -62,7 +69,7 @@ const parserWithMetadata = parser.configure({
       ScopeDeclaration: indentDeclaration,
       PredicateBody: indentPredicateOrRule,
       // FunctionBody: indentPredicateOrRule,
-      // RuleBody: indentPredicateOrRule,
+      RuleBody: indentPredicateOrRule,
       BlockComment: indentBlockComment,
     }),
     foldNodeProp.add({
@@ -71,9 +78,9 @@ const parserWithMetadata = parser.configure({
       ParameterList: foldInside,
       PredicateBody: foldInside,
       // FunctionBody: foldInside,
-      // RuleBody: foldInside,
+      RuleBody: foldInside,
       Conjunction: foldConjunction,
-      // Consequent: foldWholeNode,
+      Consequent: foldWholeNode,
       AtomDeclaration: foldDeclaration,
       NodeDeclaration: foldDeclaration,
       ScopeDeclaration: foldDeclaration,
@@ -92,7 +99,7 @@ const problemLanguage = LRLanguage.define({
       },
       line: '%',
     },
-    indentOnInput: /^\s*(?:\{|\}|\(|\)|->|;|\.)$/,
+    indentOnInput: /^\s*(?:\{|\}|\(|\)|->|==>|;|\.)$/,
   },
 });