WIP Improved cover analysis output

tla-assignments/src/main/scala/at/forsyte/apalache/tla/assignments/CoverChecker.scala

@@ -1,8 +1,8 @@
 package at.forsyte.apalache.tla.assignments
 
-import at.forsyte.apalache.tla.assignments.CoverData.ProblemData
 import at.forsyte.apalache.tla.lir.oper._
 import at.forsyte.apalache.tla.lir._
+import at.forsyte.apalache.tla.lir.storage.BodyMap
 
 import scala.collection.immutable.{Map, Set}
 
@@ -13,97 +13,126 @@ import scala.collection.immutable.{Map, Set}
   */
 class CoverChecker( allVariables: Set[String], manuallyAssigned: Set[String] = Set.empty ) {
 
-  type letInOperBodyMapType = Map[String, CoverData]
 
-  private def mkCoverInternal( initialLetInOperBodyMap : letInOperBodyMapType )
-                           ( ex: TlaEx ): CoverData = ex match {
+  type operCoverMapType = Map[String, CoverData]
+
+  def mkCover( bodyMap: BodyMap ) : operCoverMapType = bodyMap.values.foldLeft( Map.empty[String,CoverData] ){
+    case (partialMap, TlaOperDecl( opName, _, opBody)) =>
+      val (opCover, sideComputedMap) = mkCoverInternal( bodyMap, partialMap, Map.empty )( opBody )
+      sideComputedMap + ( opName -> opCover )
+  }
+
+  def mkCoverEx( bodyMap : BodyMap )( ex : TlaEx, initialCoverMap : operCoverMapType = Map.empty ) : CoverData =
+    mkCoverInternal( bodyMap, initialCoverMap, Map.empty )( ex )._1
+
+  private def mkCoverInternal(
+                               bodyMap : BodyMap,
+                               initialOperMap : operCoverMapType,
+                               initialLetInOperBodyMap : operCoverMapType
+                             )
+                             ( ex: TlaEx ): (CoverData,operCoverMapType)  = ex match {
     /** Recursive case, connectives */
     case OperEx( oper, args@_* ) if oper == TlaBoolOper.and || oper == TlaBoolOper.or =>
 
       /** First, process children */
-      val processedChildArgs : Seq[CoverData] =
-        args.map( mkCoverInternal(initialLetInOperBodyMap) )
+      val processedChildArgs : Seq[(CoverData, operCoverMapType)] =
+        args.map( mkCoverInternal( bodyMap, initialOperMap, initialLetInOperBodyMap ) )
 
-      /** Compute parent cover from children */
+      val (childCover,newSideComputedMap) = processedChildArgs.foldLeft( Seq.empty[CoverData], initialOperMap ){
+        case ((partialChildren, partialMap), (coverData, map) ) =>
+          (partialChildren :+ coverData, partialMap ++ map)
+      }
 
-      if ( oper == TlaBoolOper.and )
-        NonBranch( ex.ID, processedChildArgs:_* ) else
-        BranchPoint( ex.ID, processedChildArgs : _* )
+      /** Compute parent cover from children */
+      val parentCover = if ( oper == TlaBoolOper.and )
+        NonBranch( ex.ID, childCover:_* ) else
+        BranchPoint( ex.ID, childCover : _* )
 
+      (parentCover, newSideComputedMap)
 
     /** Base case, assignment candidates */
-    case OperEx( TlaOper.eq, OperEx( TlaActionOper.prime, NameEx( name ) ), star ) =>
+    case OperEx( TlaOper.eq, OperEx( TlaActionOper.prime, NameEx( name ) ), _ ) =>
       /** it's a candidate for name iff name \notin manuallyAssigned */
-      if ( !manuallyAssigned.contains( name ) ) Candidate( name, ex.ID ) else NonCandidate( ex.ID )
-    case OperEx( BmcOper.assign, OperEx( TlaActionOper.prime, NameEx( name ) ), star ) =>
+      val cover = if ( !manuallyAssigned.contains( name ) ) Candidate( name, ex.ID ) else NonCandidate( ex.ID )
+      (cover, initialOperMap)
+
+    case OperEx( BmcOper.assign, OperEx( TlaActionOper.prime, NameEx( name ) ), _ ) =>
       /** it's a candidate for name iff name \in manuallyAssigned */
-      if ( manuallyAssigned.contains( name ) ) Candidate( name, ex.ID ) else NonCandidate( ex.ID )
+      val cover = if ( manuallyAssigned.contains( name ) ) Candidate( name, ex.ID ) else NonCandidate( ex.ID )
+      (cover, initialOperMap)
 
     /** Recursive case, quantifier */
-    case OperEx( TlaBoolOper.exists, NameEx( _ ), star, subEx ) =>
-      mkCoverInternal( initialLetInOperBodyMap )( subEx )
+    case OperEx( TlaBoolOper.exists, NameEx( _ ), _, subEx ) =>
+      mkCoverInternal( bodyMap, initialOperMap, initialLetInOperBodyMap )( subEx )
 
-    case OperEx( TlaControlOper.ifThenElse, star, thenExpr, elseExpr ) =>
+    case OperEx( TlaControlOper.ifThenElse, _, thenExpr, elseExpr ) =>
       /** Recurse on both branches */
-      val thenResults = mkCoverInternal( initialLetInOperBodyMap )(thenExpr)
-      val elseResults = mkCoverInternal( initialLetInOperBodyMap )(elseExpr)
+      val (thenCover, thenMap) = mkCoverInternal( bodyMap, initialOperMap, initialLetInOperBodyMap )( thenExpr )
+      val (elseCover, elseMap) = mkCoverInternal( bodyMap, initialOperMap, initialLetInOperBodyMap )( elseExpr )
 
       /** Continue as with disjunction */
-      BranchPoint( ex.ID,  thenResults, elseResults )
+      val cover = BranchPoint( ex.ID,  thenCover, elseCover )
+      (cover, thenMap  ++ elseMap)
 
     /** Recursive case, nullary LetIn */
     case LetInEx( body, defs@_* ) =>
-      // Sanity check, all operators must be nullary
-      assert( defs.forall { _.formalParams.isEmpty } )
       /** First, analyze the bodies, to reuse later */
-      val bodyResults = (defs map { d =>
-        d.name -> mkCoverInternal( initialLetInOperBodyMap )( d.body )
-      }).toMap
+      val (newLetInMap, sideComputedMap) = defs.foldLeft( initialLetInOperBodyMap, initialOperMap ) {
+        case ((partialLetInMap, partialSideComputation), TlaOperDecl( letInDeclName, _, letInDeclbody )) =>
+          val (cover, newSideComputedMap) =
+            mkCoverInternal( bodyMap, partialSideComputation, partialLetInMap )( letInDeclbody )
+          (partialLetInMap + ( letInDeclName -> cover ), newSideComputedMap)
+      }
 
       /** Then, analyze the body, with the bodyResults map */
-      mkCoverInternal( initialLetInOperBodyMap ++ bodyResults )( body )
+      mkCoverInternal( bodyMap, sideComputedMap, newLetInMap )( body )
+
+    /** Application case, user-defined top-level operator */
+    case OperEx( TlaOper.apply, NameEx(operName), _* ) if bodyMap.contains( operName ) =>
+      // Assignments cannot appear in arguments to an operator, so we can ignore those
+
+      // We can sometimes leverage memoization:
+      val preComputedOpt = initialOperMap.get( operName )
+      if (preComputedOpt.nonEmpty)
+        (preComputedOpt.get, initialOperMap)
+      else {
+        // If not, we compute the cover and store it
+        val opBody = bodyMap( operName ).body
+        val (cover, sideComputedMap) = mkCoverInternal( bodyMap, initialOperMap, initialLetInOperBodyMap )( opBody )
+        (cover, sideComputedMap + (operName -> cover))
+      }
 
-    /** Nullary apply */
+
+    /** Other apply */
     case OperEx( TlaOper.apply, NameEx(operName) ) =>
       // Apply may appear in higher order operators, so it might not be possible to pre-analyze
-      initialLetInOperBodyMap.getOrElse( operName, NonCandidate( ex.ID ) )
+      val cover = initialLetInOperBodyMap.getOrElse( operName, NonCandidate( ex.ID ) )
+      (cover, initialOperMap)
 
     /** In the other cases, return the default args */
-    case _ => NonCandidate( ex.ID )
+    case _ => (NonCandidate( ex.ID ), initialOperMap)
   }
 
-  def mkCover( ex: TlaEx) : CoverData = {
-    mkCoverInternal( Map.empty )(ex)
-  }
 
-  /** Computes the set of all variables, which are covered by `ex` */
-  def coveredVars( ex: TlaEx ) : Set[String] = {
-    val cd = mkCover(ex)
-
-    val potentialProblemMap = (allVariables map {
-      v => v -> CoverData.uncoveredBranchPoints( v )( cd )
-    }).toMap
-
-
-    potentialProblemMap.keySet.filter { k => potentialProblemMap(k).noProblem }
+  /** Computes the set of all variables, which are covered by each operator */
+  def coveredVars( coverMap: operCoverMapType ) : Map[ String, Set[String]] = {
+    coverMap map {
+      case (opName, cover) => opName ->
+        allVariables.filter(
+          v => CoverData.uncoveredBranchPoints( v )( cover ).noProblem
+        )
+    }
 
   }
 
-  def findProblems( ex: TlaEx ): Option[Map[ String, ProblemData ]] = {
-    val cd = mkCover(ex)
-
-    val potentialProblemMap = (allVariables map {
-      v => v -> CoverData.uncoveredBranchPoints( v )( cd )
-    }).toMap
-
-
-    val problemMap = potentialProblemMap.filter{ case (k,v) => !v.noProblem }
-
-    if (problemMap.isEmpty) {
-      None
-    } else {
-      Some( problemMap )
-    }
+  /** Computes a set of assignment candidates under each operator */
+  def assignmentLocaitons( coverMap: operCoverMapType ) : Map[String, Set[Candidate]] =
+    coverMap map {
+      case (opName, cover) => opName ->
+        allVariables.flatMap(
+          v => CoverData.uncoveredBranchPoints( v )( cover ).problemUIDs.asgns
+        )
 
   }
+
 }

tla-assignments/src/main/scala/at/forsyte/apalache/tla/assignments/CoverData.scala

@@ -17,13 +17,14 @@ object CoverData{
 
   class CoverException( s : String ) extends Exception( s )
 
-  class IncompleteCover(seq: Seq[UID]) {
+  class IncompleteCover( seq: Seq[UID], assignments: Seq[Candidate]) {
     def get: Seq[UID] = seq
+    def asgns: Seq[Candidate] = assignments
     def isEmpty: Boolean = seq.isEmpty
   }
-  sealed case class AtLeastOne( seq: Seq[UID] ) extends IncompleteCover(seq)
-  sealed case class All( seq: Seq[UID] ) extends IncompleteCover(seq)
-  sealed case class NoProblem() extends IncompleteCover( Seq.empty )
+  sealed case class AtLeastOne( seq: Seq[UID], assignments: Seq[Candidate] ) extends IncompleteCover(seq, assignments )
+  sealed case class All( seq: Seq[UID], assignments: Seq[Candidate] ) extends IncompleteCover(seq, assignments)
+  sealed case class NoProblem( assignments: Seq[Candidate] ) extends IncompleteCover( Seq.empty, assignments)
 
   /**
     * ProblemData represents potential witnesses to cover violations.
@@ -59,22 +60,25 @@ object CoverData{
   }
 
   def uncoveredBranchPoints( varName: String )( cd: CoverData ) : ProblemData = cd match {
-    case Candidate( v, loc ) =>
+    case c@Candidate( v, loc ) =>
       // Candidate(v, _) covers varName iff v == varName
-      val problem = if ( varName == v ) NoProblem() else AtLeastOne( Seq( loc ) )
+      val problem = if ( varName == v ) NoProblem( Seq( c ) ) else AtLeastOne( Seq( loc ), Seq.empty )
       ProblemData( problem, Map.empty )
     case NonCandidate( loc ) =>
       // NonCandidate never covers varname, but is also a leaf in blameMap
-      ProblemData( AtLeastOne( Seq( loc ) ), Map.empty )
+      ProblemData( AtLeastOne( Seq( loc ), Seq.empty ), Map.empty )
     case BranchPoint( loc, branches@_* ) =>
       // BranchPoint represents disjunction/ITE, so it covers varName iff
       // *all* disjuncts/ITE branches cover varName
       val branchIssues =
-        branches.foldLeft( ProblemData( NoProblem(), Map.empty ) ) {
+        branches.foldLeft( ProblemData( NoProblem( Seq.empty ), Map.empty ) ) {
           case (pd,brCd) =>
             val brPd = uncoveredBranchPoints( varName )( brCd )
             ProblemData(
-              All( pd.problemUIDs.get ++ brPd.problemUIDs.get ),
+              All(
+                pd.problemUIDs.get ++ brPd.problemUIDs.get,
+                pd.problemUIDs.asgns ++ brPd.problemUIDs.asgns
+              ),
               pd.blameMap ++ brPd.blameMap
             )
         }
@@ -83,7 +87,7 @@ object CoverData{
         // If a subexpression contains a cover violation the BranchPoint sets itself
         // as a problem location and pushes the subexpression issues into blameMap
         ProblemData(
-          All( Seq( loc ) ),
+          All( Seq( loc ), branchIssues.problemUIDs.asgns ),
           branchIssues.blameMap + ( loc -> branchIssues.problemUIDs )
         )
       } else {
@@ -94,21 +98,25 @@ object CoverData{
       // NonBranch corresponds to conjunction, so it covers varName if any of its
       // subexpressions cover varName
       val elemIssues = elements map uncoveredBranchPoints( varName )
-      if ( elemIssues.exists( _.problemUIDs.isEmpty ) ) {
+      val coverOpt = elemIssues.find( _.problemUIDs.isEmpty )
+      if ( coverOpt.nonEmpty ) {
         // If a suitable cover is found, report NoProblem
-        ProblemData( NoProblem(), Map.empty )
+        ProblemData( NoProblem( coverOpt.get.problemUIDs.asgns ), Map.empty )
       } else {
         // Otherwise, no sub-expression covers varName, so we aggregate all sub-locations
-        val problemAggregate = elemIssues.foldLeft( ProblemData( NoProblem(), Map.empty ) ) {
+        val problemAggregate = elemIssues.foldLeft( ProblemData( NoProblem(Seq.empty), Map.empty ) ) {
           case (pd,brPd) =>
             ProblemData(
-              AtLeastOne( pd.problemUIDs.get ++ brPd.problemUIDs.get),
+              AtLeastOne(
+                pd.problemUIDs.get ++ brPd.problemUIDs.get,
+                pd.problemUIDs.asgns ++ brPd.problemUIDs.asgns
+              ),
               pd.blameMap ++ brPd.blameMap
             )
         }
         // Then, NonBranch sets itself as a problem location, as in the previous case
         ProblemData(
-          AtLeastOne( Seq( loc ) ),
+          AtLeastOne( Seq( loc ), problemAggregate.problemUIDs.asgns ),
           problemAggregate.blameMap + ( loc -> problemAggregate.problemUIDs )
         )
       }