Allow adt constructors to have associated axiom

src/main/scala/viper/silver/parser/ParseAst.scala

@@ -268,6 +268,11 @@ sealed trait PIdnUseName[T <: PDeclarationInner] extends PIdnUse {
 /** Any `PNode` which should be ignored (as well as it's children) by the `NameAnalyser`. */
 trait PNameAnalyserOpaque extends PNode
 
+trait PNameAnalyserCustom extends PNode {
+  def nameDown(map: NameAnalyserCtxt): Unit
+  def nameUp(map: NameAnalyserCtxt): Unit
+}
+
 case class PIdnRef[T <: PDeclarationInner](name: String)(val pos: (Position, Position))(implicit val ctag: scala.reflect.ClassTag[T]) extends PIdnUseName[T] {
   override def rename(newName: String): PIdnUse = PIdnRef(newName)(pos)
   /** Changes the type of declaration which is referenced, preserves all previously added `decls` but discards `filters`. */
@@ -1691,7 +1696,11 @@ case class PDomainFunction(annotations: Seq[PAnnotation], unique: Option[PKw.Uni
   override def endLinebreak = false
 }
 
-case class PAxiom(annotations: Seq[PAnnotation], axiom: PKw.Axiom, idndef: Option[PIdnDef], exp: PBracedExp)(val pos: (Position, Position)) extends PDomainMember
+/** Declares that all children expressions must always be well-defined. This is
+ * enforced during type-checking. Used e.g. for axioms. */
+trait PAlwaysWellDefined
+
+case class PAxiom(annotations: Seq[PAnnotation], axiom: PKw.Axiom, idndef: Option[PIdnDef], exp: PBracedExp)(val pos: (Position, Position)) extends PDomainMember with PAlwaysWellDefined
 
 case class PDomainInterpretation(name: PRawString, c: PSym.Colon, lit: PStringLiteral)(val pos: (Position, Position)) extends PNode
 case class PDomainInterpretations(k: PReserved[PKeywordLang], m: PDelimited.Comma[PSym.Paren, PDomainInterpretation])(val pos: (Position, Position)) extends PNode {

src/main/scala/viper/silver/parser/Resolver.scala

@@ -699,7 +699,7 @@ case class TypeChecker(names: NameAnalyser) {
                           check(fd.typ)
                           fd.formalArgs foreach (a => check(a.typ))
                         }
-                        if (pfa.isDescendant[PAxiom] && pfn.pres.toSeq.exists(pre => pre.k.rs == Requires)) {
+                        if (pfa.isDescendant[PAlwaysWellDefined] && pfn.pres.toSeq.exists(pre => pre.k.rs == Requires)) {
                           // A domain axiom, which must always be well-defined, is calling a function that has at least
                           // one real precondition (i.e., not just a requires clause or something similar that's
                           // temporarily represented as a precondition), which means that the call may not always be
@@ -941,6 +941,13 @@ case class DeclarationMap(
   def keys = decls.keys
 }
 
+// Used by `PNameAnalyserCustom`
+trait NameAnalyserCtxt {
+  def getMap(): DeclarationMap
+  def pushScope(s: PScope): Unit
+  def popScope(): Unit
+}
+
 /**
   * Resolves identifiers to their declaration. The important traits that relate to this are:
   * - `PDeclaration` marks a declaration of an identifier.
@@ -960,6 +967,7 @@ case class DeclarationMap(
   * - `PGlobalUniqueDeclaration`: marks a declaration as unique within the program.
   * 
   * - `PNameAnalyserOpaque` marks a scope as opaque (should not be traversed by the name analyser).
+  * - `PNameAnalyserCustom` marks a node which implements custom logic for name analysis.
   */
 case class NameAnalyser() {
 
@@ -992,41 +1000,79 @@ case class NameAnalyser() {
 
   private val namesInScope = mutable.Set.empty[String]
 
-  def check(g: PNode, target: Option[PNode], initialCurScope: PScope = null): Unit = {
-    var curScope: PScope = initialCurScope
-    def getMap(): DeclarationMap = Option(curScope).map(_.scopeId).map(localDeclarationMaps.get(_).get).getOrElse(globalDeclarationMap)
-
+  case class NameAnalyserCtxt(names: NameAnalyser) extends viper.silver.parser.NameAnalyserCtxt {
+    var curScope: PScope = null
     val scopeStack = mutable.Stack[PScope]()
     var opaque = 0
 
+    def getMap(): DeclarationMap = Option(curScope).map(_.scopeId).map(names.localDeclarationMaps.get(_).get).getOrElse(names.globalDeclarationMap)
+
+    def pushScope(s: PScope): Unit = {
+      names.localDeclarationMaps.put(s.scopeId, DeclarationMap(isMember = s.isInstanceOf[PMember]))
+      scopeStack.push(curScope)
+      curScope = s
+    }
+
+    def popScope(): Unit = {
+      val popMap = localDeclarationMaps.get(curScope.scopeId).get
+      val newScope = scopeStack.pop()
+      curScope = newScope
+
+      val clashing = getMap().merge(popMap)
+      clashing.foreach { case (clashing, unique) =>
+        messages ++= FastMessaging.message(clashing.idndef, s"duplicate identifier `${clashing.idndef.name}` at ${clashing.idndef.pos._1} and at ${unique.idndef.pos._1}")
+      }
+    }
+
+    def finish(): Unit = {
+      // If we started from some inner scope, walk all the way back out to the whole program
+      // with a variation of nodeUpNameCollectorVisitor
+      while (curScope != null) {
+        val popMap = localDeclarationMaps.get(curScope.scopeId).get
+        curScope.getAncestor[PScope] match {
+          case Some(newScope) =>
+            curScope = newScope
+          case None =>
+            curScope = null
+        }
+        getMap().merge(popMap)
+      }
+    }
+  }
+
+  def check(g: PNode, target: Option[PNode], initialCurScope: PScope = null): Unit = {
+    val ctx: NameAnalyserCtxt = NameAnalyserCtxt(this)
+    ctx.curScope = initialCurScope
+
     val nodeDownNameCollectorVisitor = new PartialFunction[PNode, Unit] {
       def apply(n: PNode) = {
         if (n == target.orNull)
-          namesInScope ++= getMap().keys
+          namesInScope ++= ctx.getMap().keys
         n match {
           // Opaque
           case _: PNameAnalyserOpaque =>
-            opaque += 1
-          case _ if opaque > 0 =>
+            ctx.opaque += 1
+          case _ if ctx.opaque > 0 =>
+          // Custom
+          case c: PNameAnalyserCustom =>
+            c.nameDown(ctx)
           // Regular
           case d: PDeclaration =>
             // Add to declaration map
-            val localDecls = getMap()
+            val localDecls = ctx.getMap()
             localDecls.newDecl(d)
             val clashing = localDecls.checkUnique(d, true)
             if (clashing.isDefined)
               messages ++= FastMessaging.message(d.idndef, s"duplicate identifier `${d.idndef.name}` at ${d.idndef.pos._1} and at ${clashing.get.idndef.pos._1}")
           case i: PIdnUseName[_] if target.isEmpty =>
-            getMap().newRef(i)
+            ctx.getMap().newRef(i)
           case _ =>
         }
 
         n match {
-          case _ if opaque > 0 =>
-          case s: PScope =>
-            localDeclarationMaps.put(s.scopeId, DeclarationMap(isMember = s.isInstanceOf[PMember]))
-            scopeStack.push(curScope)
-            curScope = s
+          case _ if ctx.opaque > 0 =>
+          case _: PNameAnalyserCustom =>
+          case s: PScope => ctx.pushScope(s)
           case _ =>
         }
       }
@@ -1037,6 +1083,7 @@ case class NameAnalyser() {
           case _: PScope => true
           case _: PIdnUseName[_] => true
           case _: PNameAnalyserOpaque => true
+          case _: PNameAnalyserCustom => true
           case _ => target.isDefined
         }
       }
@@ -1047,18 +1094,13 @@ case class NameAnalyser() {
         n match {
           // Opaque
           case _: PNameAnalyserOpaque =>
-            opaque -= 1
-          case _ if opaque > 0 =>
+            ctx.opaque -= 1
+          case _ if ctx.opaque > 0 =>
+          // Custom
+          case c: PNameAnalyserCustom =>
+            c.nameUp(ctx)
           // Regular
-          case _: PScope =>
-            val popMap = localDeclarationMaps.get(curScope.scopeId).get
-            val newScope = scopeStack.pop()
-            curScope = newScope
-
-            val clashing = getMap().merge(popMap)
-            clashing.foreach { case (clashing, unique) =>
-              messages ++= FastMessaging.message(clashing.idndef, s"duplicate identifier `${clashing.idndef.name}` at ${clashing.idndef.pos._1} and at ${unique.idndef.pos._1}")
-            }
+          case _: PScope => ctx.popScope()
           case _ =>
         }
       }
@@ -1078,18 +1120,8 @@ case class NameAnalyser() {
     // If we started from some inner scope, walk all the way back out to the whole program
     // with a variation of nodeUpNameCollectorVisitor
     if (initialCurScope != null) {
-      assert(initialCurScope == curScope)
-
-      while (curScope != null) {
-        val popMap = localDeclarationMaps.get(curScope.scopeId).get
-        curScope.getAncestor[PScope] match {
-          case Some(newScope) =>
-            curScope = newScope
-          case None =>
-            curScope = null
-        }
-        getMap().merge(popMap)
-      }
+      assert(initialCurScope == ctx.curScope)
+      ctx.finish()
     }
   }
 

src/main/scala/viper/silver/plugin/standard/adt/AdtASTExtension.scala

@@ -66,7 +66,7 @@ case class Adt(name: String, constructors: Seq[AdtConstructor], typVars: Seq[Typ
   * @param typ        the return type of the constructor
   * @param adtName    the name corresponding of the corresponding ADT
   */
-case class AdtConstructor(name: String, formalArgs: Seq[LocalVarDecl])
+case class AdtConstructor(name: String, formalArgs: Seq[LocalVarDecl], axiom: Option[Exp])
                          (val pos: Position, val info: Info, val typ: AdtType, val adtName: String, val errT: ErrorTrafo)
   extends ExtensionMember {
 
@@ -84,19 +84,20 @@ case class AdtConstructor(name: String, formalArgs: Seq[LocalVarDecl])
     if (!forceRewrite && this.children == children && pos.isEmpty)
       this
     else {
-      assert(children.length == 2, s"AdtConstructor : expected length 2 but got ${children.length}")
+      assert(children.length == 3, s"AdtConstructor : expected length 3 but got ${children.length}")
       val first = children.head.asInstanceOf[String]
       val second = children.tail.head.asInstanceOf[Seq[LocalVarDecl]]
-      AdtConstructor(first, second)(this.pos, this.info, this.typ, this.adtName, this.errT).asInstanceOf[this.type]
+      val third = children.tail.tail.head.asInstanceOf[Option[Exp]]
+      AdtConstructor(first, second, third)(this.pos, this.info, this.typ, this.adtName, this.errT).asInstanceOf[this.type]
     }
 
   }
 }
 
 object AdtConstructor {
-  def apply(adt: Adt, name: String, formalArgs: Seq[LocalVarDecl])
+  def apply(adt: Adt, name: String, formalArgs: Seq[LocalVarDecl], axiom: Option[Exp])
            (pos: Position = NoPosition, info: Info = NoInfo, errT: ErrorTrafo = NoTrafos): AdtConstructor = {
-    AdtConstructor(name, formalArgs)(pos, info, AdtType(adt, (adt.typVars zip adt.typVars).toMap), adt.name, errT)
+    AdtConstructor(name, formalArgs, axiom)(pos, info, AdtType(adt, (adt.typVars zip adt.typVars).toMap), adt.name, errT)
   }
 }
 

src/main/scala/viper/silver/plugin/standard/adt/AdtPASTExtension.scala

@@ -108,7 +108,7 @@ object PAdtFieldDecl {
   def apply(d: PIdnTypeBinding): PAdtFieldDecl = PAdtFieldDecl(d.idndef, d.c, d.typ)(d.pos)
 }
 
-case class PAdtConstructor(annotations: Seq[PAnnotation], idndef: PIdnDef, args: PDelimited.Comma[PSym.Paren, PAdtFieldDecl])(val pos: (Position, Position)) extends PExtender with PNoSpecsFunction with PGlobalUniqueDeclaration with PAdtChild {
+case class PAdtConstructor(annotations: Seq[PAnnotation], idndef: PIdnDef, args: PDelimited.Comma[PSym.Paren, PAdtFieldDecl], axiom: Option[PAdtConstructorAxiom])(val pos: (Position, Position)) extends PExtender with PNoSpecsFunction with PGlobalUniqueDeclaration with PAdtChild {
   override def resultType: PType = adt.getAdtType
   def fieldDecls: Seq[PAdtFieldDecl] = this.args.inner.toSeq
   override def typecheck(t: TypeChecker, n: NameAnalyser): Option[Seq[String]] = {
@@ -119,14 +119,16 @@ case class PAdtConstructor(annotations: Seq[PAnnotation], idndef: PIdnDef, args:
       val decl = discClashes.head
       t.messages ++= FastMessaging.message(idndef, "corresponding adt discriminator identifier `" + decl.idndef.name + "` at " + idndef.pos._1 + " is shadowed at " + decl.idndef.pos._1)
     }
+    axiom foreach (_.typecheck(t, n))
     None
   }
 
   override def translateMemberSignature(t: Translator): AdtConstructor = {
     AdtConstructor(
       PAdt.findAdt(adt.idndef, t),
       idndef.name,
-      formalArgs map (_.asInstanceOf[PAdtFieldDecl]) map { arg => LocalVarDecl(arg.idndef.name, t.ttyp(arg.typ))(t.liftPos(arg.idndef)) }
+      formalArgs map (_.asInstanceOf[PAdtFieldDecl]) map { arg => LocalVarDecl(arg.idndef.name, t.ttyp(arg.typ))(t.liftPos(arg.idndef)) },
+      axiom map (axiom => t.exp(axiom.e.inner))
     )(t.liftPos(this), Translator.toInfo(annotations, this))
   }
 
@@ -150,6 +152,24 @@ object PAdtConstructor {
   def findAdtConstructor(id: PIdentifier, t: Translator): AdtConstructor = t.getMembers()(id.name).asInstanceOf[AdtConstructor]
 }
 
+case class PAdtConstructorAxiom(axiom: PKw.Axiom, e: PGrouped.Paren[PExp])(val pos: (Position, Position)) extends PExtender with PDomainMember with PNameAnalyserCustom with PAlwaysWellDefined {
+  override def nameDown(ctx: NameAnalyserCtxt): Unit = {
+    ctx.pushScope(this)
+    val map = ctx.getMap()
+    val args = getAncestor[PAdtConstructor].get.args.inner.toSeq
+    args.map(a => PLogicalVarDecl(a.idndef, a.c, a.typ)(a.pos)).foreach(map.newDecl)
+  }
+  override def nameUp(ctx: NameAnalyserCtxt): Unit = ctx.popScope()
+
+  // We extend `PDomainMember` to make this a `PScope`
+  override def domain: PDomain = null
+
+  override def typecheck(t: TypeChecker, n: NameAnalyser): Option[Seq[String]] = {
+    t.check(e.inner, TypeHelper.Bool)
+    None
+  }
+}
+
 case class PAdtDeriving(k: PReserved[PDerivesKeyword.type], derivingInfos: PAdtSeq[PAdtDerivingInfo])(val pos: (Position, Position)) extends PExtender with PAdtChild {
   override def typecheck(t: TypeChecker, n: NameAnalyser): Option[Seq[String]] = {
     derivingInfos.inner foreach (_.typecheck(t, n))

src/main/scala/viper/silver/plugin/standard/adt/AdtPlugin.scala

@@ -20,7 +20,7 @@ class AdtPlugin(@unused reporter: viper.silver.reporter.Reporter,
                 config: viper.silver.frontend.SilFrontendConfig,
                 fp: FastParser) extends SilverPlugin with ParserPluginTemplate {
 
-  import fp.{annotation, argList, idnTypeBinding, idndef, idnref, semiSeparated, typ, typeList, domainTypeVarDecl, ParserExtension, lineCol, _file}
+  import fp.{annotation, argList, idnTypeBinding, idndef, idnref, parenthesizedExp, semiSeparated, typ, typeList, domainTypeVarDecl, ParserExtension, lineCol, _file}
   import FastParserCompanion.{ExtendedParsing, LeadingWhitespace, PositionParsing, reservedKw, reservedSym, whitespace}
 
   /**
@@ -91,7 +91,9 @@ class AdtPlugin(@unused reporter: viper.silver.reporter.Reporter,
 
   def adtConstructorDecls[$: P]: P[PAdtSeq[PAdtConstructor]] = P(semiSeparated(adtConstructorDecl).braces.map(PAdtSeq.apply _)).pos
 
-  def adtConstructorDecl[$: P]: P[PAdtConstructor] = P((annotation.rep ~ idndef ~ argList(idnTypeBinding.map(PAdtFieldDecl(_)))) map (PAdtConstructor.apply _).tupled).pos
+  def adtConstructorDecl[$: P]: P[PAdtConstructor] = P((annotation.rep ~ idndef ~ argList(idnTypeBinding.map(PAdtFieldDecl(_))) ~ adtConstructorAxiom.?) map (PAdtConstructor.apply _).tupled).pos
+
+  def adtConstructorAxiom[$: P]: P[PAdtConstructorAxiom] = P((P(PKw.Axiom) ~ parenthesizedExp) map (PAdtConstructorAxiom.apply _).tupled).pos
 
   override def beforeResolve(input: PProgram): PProgram = {
     if (deactivated) {