Use a Map from Snap to Snap to represent a magic wand snapshot.

src/main/resources/magic_wand_snap_functions_declarations.smt2

@@ -0,0 +1 @@
+(declare-fun MWSF_apply ($MWSF $Snap) $Snap)

src/main/scala/Utils.scala

@@ -222,6 +222,7 @@ package object utils {
     case class ViperEmbedding(embeddedSort: Sort) extends silver.ast.ExtensionType {
       def substitute(typVarsMap: Predef.Map[silver.ast.TypeVar, silver.ast.Type]): silver.ast.Type = this
       def isConcrete: Boolean = true
+      override def toString: String = s"ViperEmbedding(sorts.$embeddedSort)"
     }
   }
 

src/main/scala/decider/TermToSMTLib2Converter.scala

@@ -61,6 +61,8 @@ class TermToSMTLib2Converter
 
     case sorts.FieldPermFunction() => text("$FPM")
     case sorts.PredicatePermFunction() => text("$PPM")
+
+    case sorts.MagicWandSnapFunction => text("$MWSF")
   }
 
   def convert(d: Decl): String = {
@@ -261,15 +263,16 @@ class TermToSMTLib2Converter
 
     case Domain(id, fvf) => parens(text("$FVF.domain_") <> id <+> render(fvf))
 
-    case Lookup(field, fvf, at) => //fvf.sort match {
-//      case _: sorts.PartialFieldValueFunction =>
-        parens(text("$FVF.lookup_") <> field <+> render(fvf) <+> render(at))
-//      case _: sorts.TotalFieldValueFunction =>
-//        render(Apply(fvf, Seq(at)))
-//        parens("$FVF.lookup_" <> field <+> render(fvf) <+> render(at))
-//      case _ =>
-//        sys.error(s"Unexpected sort '${fvf.sort}' of field value function '$fvf' in lookup term '$term'")
-//    }
+    case Lookup(field, fvf, at) =>
+    // fvf.sort match {
+    //   case _: sorts.PartialFieldValueFunction =>
+      parens(text("$FVF.lookup_") <> field <+> render(fvf) <+> render(at))
+    //   case _: sorts.TotalFieldValueFunction =>
+    //     render(Apply(fvf, Seq(at)))
+    //     parens("$FVF.lookup_" <> field <+> render(fvf) <+> render(at))
+    //   case _ =>
+    //     sys.error(s"Unexpected sort '${fvf.sort}' of field value function '$fvf' in lookup term '$term'")
+    // }
 
     case FieldTrigger(field, fvf, at) => parens(text("$FVF.loc_") <> field <+> (fvf.sort match {
       case sorts.FieldValueFunction(_, _) => render(Lookup(field, fvf, at)) <+> render(at)
@@ -313,6 +316,10 @@ class TermToSMTLib2Converter
       val docBindings = ssep((bindings.toSeq map (p => parens(render(p._1) <+> render(p._2)))).to(collection.immutable.Seq), space)
       parens(text("let") <+> parens(docBindings) <+> render(body))
 
+    case snap: MagicWandSnapSingleton => renderBinaryOp("$Snap.combine", snap)
+    case MagicWandSnapFunction(_, _, wandMap) => render(wandMap)
+    case MWSFLookup(mwsf, snap) => renderApp("MWSF_apply", Seq(mwsf, snap), sorts.Snap)
+
     case _: MagicWandChunkTerm
        | _: Quantification =>
       sys.error(s"Unexpected term $term cannot be translated to SMTLib code")

src/main/scala/decider/TermToZ3APIConverter.scala

@@ -435,6 +435,8 @@ class TermToZ3APIConverter
 
       case bop: Combine =>
         ctx.mkApp(combineConstructor, convertTerm(bop.p0), convertTerm(bop.p1))
+      case MagicWandSnapSingleton(abstractLhs, rhsSnapshot) =>
+        ctx.mkApp(combineConstructor, convertTerm(abstractLhs), convertTerm(rhsSnapshot))
 
       case SortWrapper(t, to) =>
         createApp(convertId(SortWrapperId(t.sort, to)), Seq(t), to)

src/main/scala/rules/Producer.scala

@@ -397,8 +397,27 @@ object producer extends ProductionRules {
           Q(s2, v1)})
 
       case wand: ast.MagicWand =>
-        val snap = sf(sorts.Snap, v)
-        magicWandSupporter.createChunk(s, wand, MagicWandSnapshot(snap), pve, v)((s1, chWand, v1) =>
+        val magicWandSnapshot: MagicWandSnapshot = if (magicWandSupporter.useMWSF(s.program)) {
+          val snapRhs = sf(sorts.MagicWandSnapFunction, v)
+
+          // Create Map that takes a snapshot, which represent the values of the consumed LHS of the wand,
+          // and relates it to the snapshot of the RHS. We use this to preserve values of the LHS in the RHS snapshot.
+          v.decider.prover.comment(s"Produce new magic wand snapshot map for wand $wand")
+          val abstractLhs = v.decider.fresh(sorts.Snap)
+          val wandMap = v.decider.fresh("$mwsf", sorts.MagicWandSnapFunction)
+          val snapFunction = MagicWandSnapFunction(abstractLhs, MWSFLookup(snapRhs, abstractLhs), wandMap)
+
+          // We assume that the wandMap that we get from `snapRhs`, which potentially is nested in a binary tree,
+          // is the same as our newly created wandMap.
+          v.decider.assumeDefinition(snapFunction.definition)
+          snapFunction
+
+        } else {
+          val snap = sf(sorts.Snap, v)
+          MagicWandSnapSingleton(snap)
+        }
+
+        magicWandSupporter.createChunk(s, wand, magicWandSnapshot, pve, v)((s1, chWand, v1) =>
           chunkSupporter.produce(s1, s1.h, chWand, v1)((s2, h2, v2) =>
             Q(s2.copy(h = h2), v2)))
 

src/main/scala/rules/QuantifiedChunkSupport.scala

@@ -634,7 +634,6 @@ object quantifiedChunkSupporter extends QuantifiedChunkSupport {
 
   /* Snapshots */
 
-  /** @inheritdoc */
   def singletonSnapshotMap(s: State,
                            resource: ast.Resource,
                            arguments: Seq[Term],

src/main/scala/state/Chunks.scala

@@ -176,13 +176,16 @@ case class MagicWandChunk(id: MagicWandIdentifier,
   override val resourceID = MagicWandID
 
   override def withPerm(newPerm: Term) = MagicWandChunk(id, bindings, args, snap, newPerm)
-  override def withSnap(newSnap: Term) = MagicWandChunk(id, bindings, args, MagicWandSnapshot(newSnap), perm)
+  override def withSnap(newSnap: Term) = newSnap match {
+    case s: MagicWandSnapshot => MagicWandChunk(id, bindings, args, s, perm)
+    case _ => sys.error(s"MagicWand snapshot has to be of type MagicWandSnapshot but found ${newSnap.getClass}")
+  }
 
   override lazy val toString = {
     val pos = id.ghostFreeWand.pos match {
       case rp: viper.silver.ast.HasLineColumn => s"${rp.line}:${rp.column}"
       case other => other.toString
     }
-    s"wand@$pos[$snap; ${args.mkString(",")}]"
+    s"wand@$pos[$snap; ${args.mkString(", ")}]"
   }
 }

src/main/scala/state/Terms.scala

@@ -76,6 +76,11 @@ object sorts {
     override lazy val toString = id.toString
   }
 
+  object MagicWandSnapFunction extends Sort {
+    val id: Identifier = Identifier("MWSF")
+    override lazy val toString: String = id.toString
+  }
+
   case class FieldPermFunction() extends Sort  {
     val id = Identifier("FPM")
     override lazy val toString = id.toString
@@ -321,39 +326,39 @@ sealed trait Term extends Node {
 
   lazy val subterms: Seq[Term] = state.utils.subterms(this)
 
-  /** @see [[ast.utility.Visitor.visit()]] */
+  /** @see [[ast.utility.Visitor.visit]] */
   def visit(f: PartialFunction[Term, Any]): Unit =
     ast.utility.Visitor.visit(this, state.utils.subterms)(f)
 
-  /** @see [[ast.utility.Visitor.visitOpt()]] */
+  /** @see [[ast.utility.Visitor.visitOpt]] */
   def visitOpt(f: Term => Boolean): Unit =
     ast.utility.Visitor.visitOpt(this, state.utils.subterms)(f)
 
-  /** @see [[ast.utility.Visitor.reduceTree()]] */
+  /** @see [[ast.utility.Visitor.reduceTree]] */
   def reduceTree[R](f: (Term, Seq[R]) => R): R =
     ast.utility.Visitor.reduceTree(this, state.utils.subterms)(f)
 
-  /** @see [[ast.utility.Visitor.existsDefined()]] */
+  /** @see [[ast.utility.Visitor.existsDefined]] */
   def existsDefined(f: PartialFunction[Term, Any]): Boolean =
     ast.utility.Visitor.existsDefined(this, state.utils.subterms)(f)
 
-  /** @see [[ast.utility.Visitor.hasSubnode()]] */
+  /** @see [[ast.utility.Visitor.hasSubnode]] */
   def hasSubterm(subterm: Term): Boolean =
     ast.utility.Visitor.hasSubnode(this, subterm, state.utils.subterms)
 
-  /** @see [[ast.utility.Visitor.deepCollect()]] */
+  /** @see [[ast.utility.Visitor.deepCollect]] */
   def deepCollect[R](f: PartialFunction[Term, R]) : Seq[R] =
     ast.utility.Visitor.deepCollect(Seq(this), state.utils.subterms)(f)
 
-  /** @see [[ast.utility.Visitor.shallowCollect()]] */
+  /** @see [[ast.utility.Visitor.shallowCollect]] */
   def shallowCollect[R](f: PartialFunction[Term, R]): Seq[R] =
     ast.utility.Visitor.shallowCollect(Seq(this), state.utils.subterms)(f)
 
-  /** @see [[ast.utility.Visitor.find()]] */
+  /** @see [[ast.utility.Visitor.find]] */
   def find[R](f: PartialFunction[Term, R]): Option[R] =
     ast.utility.Visitor.find(this, state.utils.subterms)(f)
 
-  /** @see [[state.utils.transform()]] */
+  /** @see [[state.utils.transform]] */
   def transform(pre: PartialFunction[Term, Term] = PartialFunction.empty)
                (recursive: Term => Boolean = !pre.isDefinedAt(_),
                 post: PartialFunction[Term, Term] = PartialFunction.empty)
@@ -2295,48 +2300,139 @@ object PredicateTrigger extends PreciseCondFlyweightFactory[(String, Term, Seq[T
 
 /* Magic wands */
 
-class MagicWandSnapshot(val abstractLhs: Term, val rhsSnapshot: Term) extends Combine(abstractLhs, rhsSnapshot) {
+abstract class MagicWandSnapshot(val abstractLhs: Term, val rhsSnapshot: Term) extends Term {
   utils.assertSort(abstractLhs, "abstract lhs", sorts.Snap)
-  utils.assertSort(rhsSnapshot, "rhs", sorts.Snap)
+  utils.assertSort(rhsSnapshot, "rhs snapshot", sorts.Snap)
+}
 
-  override lazy val toString = s"wandSnap(lhs = $abstractLhs, rhs = $rhsSnapshot)"
+/**
+ * Snapshot for a magic wand. It represents a function which can be used only once
+ * by equating a snapshot from consuming the wand's LHS with `abstractLhs`.
+ *
+ * For a snapshot which can be applied multiple times see [[MagicWandSnapFunction]].
+ *
+ * @param abstractLhs The abstract snapshot which acts as a placeholder in the
+ *                    `rhsSnapshot` for the values when applying the magic wand.
+ * @param rhsSnapshot Snapshot which can be used when producing the wand's RHS.
+ */
+class MagicWandSnapSingleton(override val abstractLhs: Term, override val rhsSnapshot: Term)
+  extends MagicWandSnapshot(abstractLhs, rhsSnapshot)
+    with ConditionalFlyweightBinaryOp[MagicWandSnapSingleton] {
 
-  def merge(other: MagicWandSnapshot, branchConditions: Stack[Term]): MagicWandSnapshot = {
-    assert(this.abstractLhs == other.abstractLhs)
-    val condition = And(branchConditions)
-    MagicWandSnapshot(this.abstractLhs, if (this.rhsSnapshot == other.rhsSnapshot)
-      this.rhsSnapshot
-    else
-      Ite(condition, other.rhsSnapshot, this.rhsSnapshot))
-  }
+  override lazy val toString: String = s"wandSnap(lhs = $abstractLhs, rhs = $rhsSnapshot)"
+
+  /* ConditionalFlyweightBinaryOp members */
+
+  override def sort: Sort = sorts.Snap
+
+  override def p0: Term = abstractLhs
+
+  override def p1: Term = rhsSnapshot
 }
 
-object MagicWandSnapshot  {
+object MagicWandSnapSingleton {
+  var pool = new TrieMap[(Term, Term), MagicWandSnapSingleton]()
+
+  /** Craete a new [[MagicWandSnapSingleton]] from a single snapshot term. */
   def apply(snapshot: Term): MagicWandSnapshot = {
-    assert(snapshot.sort == sorts.Snap)
+    assert(snapshot.sort == sorts.Snap, s"MagicWandSnapshot.apply expects sorts Snap but got ${snapshot.sort}")
+
     snapshot match {
-      case snap: MagicWandSnapshot => snap
-      case _ =>
-        MagicWandSnapshot(First(snapshot), Second(snapshot))
+      case snap: MagicWandSnapSingleton => snap
+      case _ => MagicWandSnapSingleton(First(snapshot), Second(snapshot))
     }
   }
 
-  // Since MagicWandSnapshot subclasses Combine, we apparently cannot inherit the normal subclass, so we
-  // have to copy paste the code here.
-  var pool = new TrieMap[(Term, Term), MagicWandSnapshot]()
+  /** Create a new [[MagicWandSnapSingleton]] from two snapshot terms. */
+  def apply(abstractLhs: Term, rhsSnapshot: Term): MagicWandSnapSingleton =
+    createIfNonExistent((abstractLhs, rhsSnapshot))
 
-  def createIfNonExistent(args: (Term, Term)): MagicWandSnapshot = {
+  /** Destruct a [[MagicWandSnapSingleton]] instance. Used in [[viper.silicon.state.utils.transform]] */
+  def unapply(mws: MagicWandSnapSingleton): Some[(Term, Term)] =
+    Some((mws.abstractLhs, mws.rhsSnapshot))
+
+  private def createIfNonExistent(args: (Term, Term)): MagicWandSnapSingleton = {
     if (Verifier.config.useFlyweight) {
       pool.getOrElseUpdate(args, actualCreate(args))
     } else {
       actualCreate(args)
     }
   }
 
-  def actualCreate(tuple: (Term, Term)) = new MagicWandSnapshot(tuple._1, tuple._2)
-  def apply(fst: Term, snd: Term): MagicWandSnapshot = createIfNonExistent((fst, snd))
+  private def actualCreate(tuple: (Term, Term)) = new MagicWandSnapSingleton(tuple._1, tuple._2)
+}
+
+/**
+ * Represents a snapshot of a magic wand, which is a map from Snap to Snap.
+ *
+ * @param abstractLhs The term that represent the snapshot of the consumed left-hand side of the magic wand.
+ *                    It is used as a bound variable in a forall quantifier.
+ * @param rhsSnapshot The term that integrates the left-hand side in the snapshot that is produced after applying the magic wand.
+ * @param wandMap The map that represents the snapshot of the magic wand. It is a variable of sort Map(Snap, Snap).
+ *                In the symbolic execution this represents a function that when we apply a magic wand
+ *                it consumes the left-hand side and uses the resulting snapshot to look up which right-hand side should be produced.
+ */
+class MagicWandSnapFunction(override val abstractLhs: Var, override val rhsSnapshot: Term, val wandMap: Term)
+  extends MagicWandSnapshot(abstractLhs, rhsSnapshot)
+    with ConditionalFlyweight[Term, MagicWandSnapFunction] {
+
+  utils.assertSort(abstractLhs, "abstract lhs", sorts.Snap)
+  utils.assertSort(rhsSnapshot, "rhs snapshot", sorts.Snap)
+  utils.assertSort(wandMap, "wand map", sorts.MagicWandSnapFunction)
+
+  override val sort: Sort = sorts.MagicWandSnapFunction
+
+  override lazy val toString = s"wandSnap($wandMap)"
+
+  override val equalityDefiningMembers: Term = wandMap
+
+  /**
+   * Creates a term that can be used to define a `wandMap`.
+   *
+   * @return Expression which says that the map applied to an arbitrary lhs equals to the snapshot of the rhs.
+   */
+  def definition: Term = Forall(
+    abstractLhs,
+    MWSFLookup(wandMap, abstractLhs) === rhsSnapshot,
+    Trigger(MWSFLookup(wandMap, abstractLhs))
+  )
+
+  /**
+   * Apply the given snapshot of the left-hand side to the magic wand map to get the snapshot of the right-hand side
+   * which includes the values of the left-hand side.
+   *
+   * @param snapLhs The snapshot of the left-hand side that should be applied to the magic wand map.
+   * @return The snapshot of the right-hand side that preserves the values of the left-hand side.
+   */
+  def applyToWandMap(snapLhs: Term): Term = MWSFLookup(wandMap, snapLhs)
+}
+
+object MagicWandSnapFunction  {
+  /** Create a new instance of [[viper.silicon.state.terms.MagicWandSnapFunction]]. */
+  def apply(abstractLhs: Var, rhsSnapshot: Term, wandMap: Term): MagicWandSnapFunction =
+    new MagicWandSnapFunction(abstractLhs, rhsSnapshot, wandMap)
 
-  def unapply(mws: MagicWandSnapshot) = Some((mws.abstractLhs, mws.rhsSnapshot))
+  /** Destructs an instance of [[viper.silicon.state.terms.MagicWandSnapFunction]]. */
+  def unapply(snap: MagicWandSnapFunction): Option[(Var, Term, Term)] =
+    Some(snap.abstractLhs, snap.rhsSnapshot, snap.wandMap)
+}
+
+class MWSFLookup(val mwsf: Term, val snap: Term) extends Term with ConditionalFlyweightBinaryOp[MWSFLookup] {
+  val sort: Sort = sorts.Snap
+  override def p0: Term = mwsf
+  override def p1: Term = snap
+  override lazy val toString = s"$mwsf[$snap]"
+}
+
+object MWSFLookup extends PreciseCondFlyweightFactory[(Term, Term), MWSFLookup] {
+  override def apply(pair: (Term, Term)): MWSFLookup = {
+    val (mwsf, snap) = pair
+    utils.assertSort(mwsf, "mwsf", sorts.MagicWandSnapFunction)
+    utils.assertSort(snap, "snap", sorts.Snap)
+    createIfNonExistent(pair)
+  }
+
+  override def actualCreate(args: (Term, Term)): MWSFLookup = new MWSFLookup(args._1, args._2)
 }
 
 class MagicWandChunkTerm(val chunk: MagicWandChunk) extends Term with ConditionalFlyweight[MagicWandChunk, MagicWandChunkTerm] {
@@ -2461,7 +2557,7 @@ object PsfTop extends (String => Identifier) {
  */
 
 /* Note: Sort wrappers should probably not be used as (outermost) triggers
- * because they are optimised away if wrappee `t` already has sort `to`.
+ * because they are optimised away if wrapped `t` already has sort `to`.
  */
 class SortWrapper(val t: Term, val to: Sort)
     extends Term

src/main/scala/state/Utils.scala

@@ -94,7 +94,7 @@ package object utils {
 
   }
 
-  /** @see [[viper.silver.ast.utility.Transformer.simplify()]] */
+  /** @see [[viper.silver.ast.utility.Simplifier.simplify]] */
   def transform[T <: Term](term: T,
                            pre: PartialFunction[Term, Term] = PartialFunction.empty)
                           (recursive: Term => Boolean = !pre.isDefinedAt(_),
@@ -196,7 +196,9 @@ package object utils {
       case MapUpdate(t0, t1, t2) => MapUpdate(go(t0), go(t1), go(t2))
       case MapDomain(t) => MapDomain(go(t))
       case MapRange(t) => MapRange(go(t))
-      case MagicWandSnapshot(lhs, rhs) => MagicWandSnapshot(go(lhs), go(rhs))
+      case MagicWandSnapSingleton(t0, t1) => MagicWandSnapSingleton(go(t0), go(t1))
+      case MagicWandSnapFunction(t0, t1, t2) => MagicWandSnapFunction(go(t0), go(t1), go(t2))
+      case MWSFLookup(t0, t1) => MWSFLookup(go(t0), go(t1))
       case Combine(t0, t1) => Combine(go(t0), go(t1))
       case First(t) => First(go(t))
       case Second(t) => Second(go(t))