Important fixes and improvements (see desc)

Support decisions with == and complex variations.
Fix wrong assumption introduced in 1f9c81
Workaround this issue: typetools/checker-framework/issues/3267

examples/comparison/iterator-attempt4/README.md

@@ -1,13 +1,6 @@
 ## Mungo Checker's output
 
 ```
-Main.java:5: error: [Object did not complete its protocol. Type: JavaIteratorProtocol{Next} | Ended] (Object did not complete its protocol. Type: JavaIteratorProtocol{Next} | Ended)
-		JavaIterator it = new JavaIterator(Arrays.asList(args).iterator());
-		             ^
-Main.java:8: error: [Cannot call next on ended protocol] (Cannot call next on ended protocol)
-      System.out.println(it.next());
-                                ^
-2 errors
 ```
 
 ## Original Mungo's output

src/main/kotlin/org/checkerframework/checker/mungo/analysis/MungoAnalysis.kt

@@ -2,6 +2,17 @@ package org.checkerframework.checker.mungo.analysis
 
 import org.checkerframework.checker.mungo.MungoChecker
 import org.checkerframework.checker.mungo.annotators.MungoAnnotatedTypeFactory
+import org.checkerframework.dataflow.analysis.Store
+import org.checkerframework.dataflow.analysis.Store.FlowRule
+import org.checkerframework.dataflow.analysis.TransferInput
+import org.checkerframework.dataflow.analysis.TransferResult
+import org.checkerframework.dataflow.cfg.block.Block
+import org.checkerframework.dataflow.cfg.block.ConditionalBlock
+import org.checkerframework.dataflow.cfg.block.RegularBlock
+import org.checkerframework.dataflow.cfg.block.SpecialBlock
+import org.checkerframework.dataflow.cfg.node.ConditionalNotNode
+import org.checkerframework.dataflow.cfg.node.Node
+import org.checkerframework.dataflow.cfg.node.ReturnNode
 import org.checkerframework.framework.flow.CFAbstractAnalysis
 import org.checkerframework.framework.flow.CFAbstractValue
 import org.checkerframework.javacutil.Pair
@@ -12,6 +23,8 @@ import javax.lang.model.type.TypeMirror
 class MungoAnalysis(checker: MungoChecker, factory: MungoAnnotatedTypeFactory, fieldValues: List<Pair<VariableElement, MungoValue>>) : CFAbstractAnalysis<MungoValue, MungoStore, MungoTransfer>(checker, factory, fieldValues) {
 
   val c = checker
+  val f = factory
+  val utils get() = c.utils
   var inClassAnalysis: Boolean = false
   var creatingInitialStore: Boolean = false
 
@@ -30,4 +43,104 @@ class MungoAnalysis(checker: MungoChecker, factory: MungoAnnotatedTypeFactory, f
       null
     } else MungoValue(this, annotations, underlyingType)
   }
+
+  fun nextIsConditional(node: Node): Boolean {
+    val block = node.block
+    if (block is RegularBlock) {
+      val succ = block.successor
+      if (succ is ConditionalBlock) {
+        return block.contents.last() === node
+      }
+    }
+    return false
+  }
+
+  private fun shouldEachToEach(node: Node?, succ: Block): Boolean {
+    return when (succ) {
+      is ConditionalBlock -> true
+      is SpecialBlock -> succ.specialType == SpecialBlock.SpecialBlockType.EXIT
+      is RegularBlock -> shouldEachToEach(node, succ.contents.firstOrNull())
+      else -> false
+    }
+  }
+
+  private fun shouldEachToEach(node: Node?, after: Node?): Boolean {
+    return when (after) {
+      is ReturnNode -> after.result === node
+      is ConditionalNotNode -> after.operand === node
+      else -> false
+    }
+  }
+
+  // We also need to control the propagation for nodes inside RegularBlock's contents
+  override fun callTransferFunction(node: Node, input: TransferInput<MungoValue, MungoStore>): TransferResult<MungoValue, MungoStore> {
+    val block = node.block
+    if (block is RegularBlock) {
+      val idx = block.contents.indexOf(node)
+      val prevIdx = idx - 1
+      if (prevIdx >= 0) {
+        val prev = block.contents[prevIdx]
+        if (!shouldEachToEach(prev, node)) {
+          return super.callTransferFunction(node, TransferInput(input.node, this, input.regularStore))
+        }
+      }
+    }
+    return super.callTransferFunction(node, input)
+  }
+
+  override fun propagateStoresTo(
+    succ: Block,
+    node: Node?,
+    currentInput: TransferInput<MungoValue, MungoStore>,
+    flowRule: FlowRule,
+    addToWorklistAgain: Boolean
+  ) {
+    when (flowRule) {
+      FlowRule.EACH_TO_EACH -> if (currentInput.containsTwoStores() && shouldEachToEach(node, succ)) {
+        addStoreBefore(
+          succ,
+          node,
+          currentInput.thenStore,
+          Store.Kind.THEN,
+          addToWorklistAgain)
+        addStoreBefore(
+          succ,
+          node,
+          currentInput.elseStore,
+          Store.Kind.ELSE,
+          addToWorklistAgain)
+      } else {
+        addStoreBefore(
+          succ,
+          node,
+          currentInput.regularStore,
+          Store.Kind.BOTH,
+          addToWorklistAgain)
+      }
+      FlowRule.THEN_TO_BOTH -> addStoreBefore(
+        succ,
+        node,
+        currentInput.thenStore,
+        Store.Kind.BOTH,
+        addToWorklistAgain)
+      FlowRule.ELSE_TO_BOTH -> addStoreBefore(
+        succ,
+        node,
+        currentInput.elseStore,
+        Store.Kind.BOTH,
+        addToWorklistAgain)
+      FlowRule.THEN_TO_THEN -> addStoreBefore(
+        succ,
+        node,
+        currentInput.thenStore,
+        Store.Kind.THEN,
+        addToWorklistAgain)
+      FlowRule.ELSE_TO_ELSE -> addStoreBefore(
+        succ,
+        node,
+        currentInput.elseStore,
+        Store.Kind.ELSE,
+        addToWorklistAgain)
+    }
+  }
 }

src/main/kotlin/org/checkerframework/checker/mungo/analysis/MungoClassAnalysis.kt

@@ -13,8 +13,11 @@ import org.checkerframework.checker.mungo.typestate.graph.DecisionState
 import org.checkerframework.checker.mungo.typestate.graph.Graph
 import org.checkerframework.checker.mungo.typestate.graph.State
 import org.checkerframework.dataflow.analysis.AnalysisResult
+import org.checkerframework.dataflow.analysis.TransferResult
 import org.checkerframework.dataflow.cfg.ControlFlowGraph
 import org.checkerframework.dataflow.cfg.UnderlyingAST
+import org.checkerframework.dataflow.cfg.node.BooleanLiteralNode
+import org.checkerframework.dataflow.cfg.node.ReturnNode
 import org.checkerframework.framework.flow.CFCFGBuilder
 import org.checkerframework.javacutil.TreeUtils
 import java.util.*
@@ -72,21 +75,23 @@ class MungoClassAnalysis(
     // The results for each method to be merged later.
     val results = mutableMapOf<MethodTree, AnalysisResult<MungoValue, MungoStore>>()
 
+    val emptyStore = initialStore.toBottom()
+
     // Update the state's store. Queue the state again if it changed.
     fun mergeStateStore(state: State, store: MungoStore) {
-      val currStore = stateToStore[state] ?: f.emptyStore
+      val currStore = stateToStore[state] ?: emptyStore
       val newStore = currStore.leastUpperBoundFields(store)
-      if (currStore != newStore) {
+      if (!stateToStore.containsKey(state) || currStore != newStore) {
         stateToStore[state] = newStore
         stateQueue.add(state)
       }
     }
 
     // Returns the merge result if it changed. Returns null otherwise.
     fun mergeMethodStore(method: MethodTree, store: MungoStore): MungoStore? {
-      val currStore = methodToStore[method] ?: f.emptyStore
+      val currStore = methodToStore[method] ?: emptyStore
       val newStore = currStore.leastUpperBoundFields(store)
-      return if (currStore != newStore) {
+      return if (!methodToStore.containsKey(method) || currStore != newStore) {
         methodToStore[method] = newStore
         newStore
       } else null
@@ -117,7 +122,10 @@ class MungoClassAnalysis(
         results[method] = a.result
 
         // Store/override exit and return statement stores
-        f.storeMethodResults(method, regularExitStore)
+        val returnStores = a.returnStatementStores
+        f.storeMethodResults(method, regularExitStore, returnStores)
+
+        val constantReturn = getConstantReturn(returnStores.map { it.first })
 
         // Merge new exit store with the stores of each destination state
         when (destState) {
@@ -126,8 +134,8 @@ class MungoClassAnalysis(
             // TODO handle enumeration values as well
             for ((label, dest) in destState.transitions) {
               when (label.label) {
-                "true" -> mergeStateStore(dest, exitResult.thenStore)
-                "false" -> mergeStateStore(dest, exitResult.elseStore)
+                "true" -> mergeStateStore(dest, if (constantReturn != false) exitResult.thenStore else emptyStore)
+                "false" -> mergeStateStore(dest, if (constantReturn != true) exitResult.elseStore else emptyStore)
                 else -> mergeStateStore(dest, regularExitStore)
               }
             }
@@ -145,4 +153,22 @@ class MungoClassAnalysis(
     f.storeClassResult(classTree, stateToStore)
   }
 
+  private fun getConstantReturn(returnStores: List<ReturnNode>): Boolean? {
+    var sawTrue = false
+    var sawFalse = false
+    for (ret in returnStores) {
+      when (val result = ret.result) {
+        is BooleanLiteralNode -> if (result.value!!) {
+          if (sawFalse) return null
+          sawTrue = true
+        } else {
+          if (sawTrue) return null
+          sawFalse = true
+        }
+        else -> return null
+      }
+    }
+    return sawTrue
+  }
+
 }