Ids in TlaDecl and refactoring of TransformationTracker (#444)

tla-bmcmt/src/main/scala/at/forsyte/apalache/tla/bmcmt/types/eager/TrivialTypeFinder.scala

@@ -98,6 +98,10 @@ class TrivialTypeFinder extends TypeFinder[CellT]
     }
   }
 
+  override def onDeclTransformation(originalDecl: TlaDecl, newDecl: TlaDecl): Unit = {
+    // ignore transformations of declarations
+  }
+
   /**
     * Given a TLA+ expression, reconstruct the types and store them in an internal storage.
     * If the expression is not well-typed, diagnostic messages can be accessed with getTypeErrors.

tla-import/src/main/scala/at/forsyte/apalache/tla/imp/SubstTranslator.scala

@@ -48,7 +48,7 @@ class SubstTranslator(sourceStore: SourceStore, context: Context) extends LazyLo
         }
 
       // copy the source info
-      sourceStore.findOrLog(ex.ID).foreach { loc => sourceStore.add(newEx.ID, loc) }
+      sourceStore.findOrWarn(ex.ID).foreach { loc => sourceStore.add(newEx.ID, loc) }
       // return
       newEx
     }

tla-import/src/main/scala/at/forsyte/apalache/tla/imp/src/SaveToStoreTracker.scala

@@ -1,7 +1,7 @@
 package at.forsyte.apalache.tla.imp.src
 
-import at.forsyte.apalache.tla.lir.TlaEx
-import at.forsyte.apalache.tla.lir.transformations.{TlaExTransformation, TransformationTracker}
+import at.forsyte.apalache.tla.lir.{TlaDecl, TlaEx, TlaOperDecl}
+import at.forsyte.apalache.tla.lir.transformations.{TlaDeclTransformation, TlaExTransformation, TransformationTracker}
 
 /**
   * This is a special form of transformation tracker that records the changes directly in the source store.
@@ -21,7 +21,37 @@ class SaveToStoreTracker(sourceStore: SourceStore) extends TransformationTracker
     input: TlaEx =>
       val output = tr(input)
       if (output.ID != input.ID) {
-        sourceStore.findOrLog(input.ID).foreach { loc => sourceStore.add(output.ID, loc) }
+        sourceStore.findOrWarn(input.ID).foreach { loc => sourceStore.add(output.ID, loc) }
+      }
+      output
+  }
+
+  /**
+    * Decorate a declaration transformation with a tracker.
+    *
+    * @param tr a declaration transformation
+    * @return a new declaration transformation that applies tr and tracks the changes
+    */
+  override def trackDecl(tr: TlaDeclTransformation): TlaDeclTransformation = {
+    input: TlaDecl =>
+      val output = tr(input)
+      if (output.ID != input.ID) {
+        sourceStore.findOrWarn(input.ID).foreach { loc => sourceStore.add(output.ID, loc) }
+      }
+      output
+  }
+
+  /**
+    * A specialized version of trackDecl, which is tuned to TlaOperDecl.
+    *
+    * @param tr a declaration transformation
+    * @return a new declaration transformation that applies tr and tracks the changes
+    */
+  override def trackOperDecl(tr: TlaOperDecl => TlaOperDecl): TlaOperDecl => TlaOperDecl = {
+    input: TlaOperDecl =>
+      val output = tr(input)
+      if (output.ID != input.ID) {
+        sourceStore.findOrWarn(input.ID).foreach { loc => sourceStore.add(output.ID, loc) }
       }
       output
   }

tla-import/src/main/scala/at/forsyte/apalache/tla/imp/src/SourceStore.scala

@@ -3,7 +3,7 @@ package at.forsyte.apalache.tla.imp.src
 import at.forsyte.apalache.tla.lir.src.{RegionTree, SourceLocation}
 import at.forsyte.apalache.tla.lir.storage.SourceMap
 import at.forsyte.apalache.tla.lir.transformations.TransformationListener
-import at.forsyte.apalache.tla.lir.{OperEx, TlaEx, UID}
+import at.forsyte.apalache.tla.lir.{OperEx, TlaDecl, TlaEx, UID}
 import com.google.inject.Singleton
 import com.typesafe.scalalogging.LazyLogging
 
@@ -59,6 +59,18 @@ class SourceStore extends TransformationListener with LazyLogging {
     }
   }
 
+  override def onDeclTransformation(originalDecl: TlaDecl, newDecl: TlaDecl): Unit = {
+    if (originalDecl.ID != newDecl.ID) {
+      find(originalDecl.ID) match {
+        case Some(loc) =>
+          add(newDecl.ID, loc)
+
+        case None =>
+          ()
+      }
+    }
+  }
+
   /**
     * Label an expression id with a source location. Unless you are sure about using this method, use addRec instead.
  *
@@ -119,7 +131,7 @@ class SourceStore extends TransformationListener with LazyLogging {
     * @param id an expression identifier
     * @return Some(location), when the expression has source information, and None otherwise
     */
-  def findOrLog(id: UID): Option[SourceLocation] = {
+  def findOrWarn(id: UID): Option[SourceLocation] = {
     find(id) match {
       case some @ Some(_) => some
       case None =>

tlair/src/main/scala/at/forsyte/apalache/tla/lir/package.scala

@@ -5,24 +5,21 @@ package lir {
   import at.forsyte.apalache.tla.lir.io.UTFPrinter
   import at.forsyte.apalache.tla.lir.oper._
 
-  /** the base class for the universe of objects used in TLA+ */
+  /** the base class for the universe of values (integers, Booleans, strings) used in TLA+ */
   abstract class TlaValue
 
   /**
     * A declaration, e.g., of a variable, constant, or an operator.
     * Technically, this class should be called TlaDef, as we are dealing with
     * TLA+ definitions, see Specifying Systems, Ch. 17.3. Unfortunately, there are
     * variable declarations and operator definitions...
-    *
-    * TODO: rename to TlaDef.
     */
-  abstract class TlaDecl extends Serializable {
+  abstract class TlaDecl extends Identifiable with Serializable {
     def name: String
+    @deprecated("Marked for removal. Use the DeepCopy transformation.")
     def deepCopy(): TlaDecl
   }
 
-  // TODO: add TlaTheoremDecl?
-
   /**
     * A module as a basic unit that contains declarations.
     *
@@ -66,21 +63,14 @@ package lir {
     override def deepCopy(): TlaAssumeDecl = TlaAssumeDecl(body.deepCopy())
   }
 
-  ///////////////// DISCUSSION
-
   /**
     * A spec, given by a list of declarations and a list of expressions.
-    *
-    * FIXME: a candidate for removal. Just use TlaModule?
     */
+  @deprecated("Marked for removal. Use TlaModule")
   case class TlaSpec( name: String, declarations: List[TlaDecl] ) extends Serializable {
     def deepCopy() : TlaSpec = TlaSpec( name, declarations.map( _.deepCopy() ) )
-
   }
 
-  ///////////////// END of DISCUSSION
-
-
   /**
   A formal parameter of an operator.
     */
@@ -107,6 +97,7 @@ package lir {
 
     def toSimpleString: String = ""
 
+    @deprecated("Marked for removal. Use the DeepCopy transformation.")
     def deepCopy() : TlaEx
   }
 

tlair/src/main/scala/at/forsyte/apalache/tla/lir/storage/ChangeListener.scala

@@ -1,8 +1,7 @@
 package at.forsyte.apalache.tla.lir.storage
 
-
 import at.forsyte.apalache.tla.lir.transformations.TransformationListener
-import at.forsyte.apalache.tla.lir.{LetInEx, OperEx, TlaEx, UID}
+import at.forsyte.apalache.tla.lir.{LetInEx, OperEx, TlaDecl, TlaEx, UID}
 import com.google.inject.Singleton
 
 import scala.collection.mutable
@@ -28,49 +27,60 @@ class ChangeListener extends TransformationListener {
 
   private def conditionalInsert(ex: TlaEx, uid: UID): Unit =
     if (!map.contains(ex.ID) && ex.ID != uid)
-      map.update( ex.ID, uid )
+      map.update(ex.ID, uid)
 
   // Sometimes, a transformation will create a complex expression from a simple one,
   // e.g. UNCHANGED x -> x' \in {x}
   // In these cases, we treat every subexpression of the resulting complex expression
   // as if it were produced by an anonymous transformation from the original
-  private def inheritUID(ex: TlaEx, uid: UID) : Unit = {
+  private def inheritUID(ex: TlaEx, uid: UID): Unit = {
     ex match {
-      case LetInEx( body, defs@_* ) =>
+      case LetInEx(body, defs @ _*) =>
         defs map {
           _.body
         } foreach {
-          inheritUID( _, uid )
+          inheritUID(_, uid)
         }
-        inheritUID( body, uid )
+        inheritUID(body, uid)
 
-      case OperEx( _, args@_* ) =>
+      case OperEx(_, args @ _*) =>
         args foreach {
-          inheritUID( _, uid )
+          inheritUID(_, uid)
         }
       case _ => // noop
     }
-    conditionalInsert( ex, uid )
+    conditionalInsert(ex, uid)
+  }
+
+  override def onTransformation(originEx: TlaEx, newEx: TlaEx): Unit = {
+    if (originEx.ID != newEx.ID) {
+      map.update(newEx.ID, originEx.ID)
+      inheritUID(newEx, originEx.ID)
+    }
   }
 
-  override def onTransformation( originEx : TlaEx, newEx : TlaEx ) : Unit = {
-    if ( originEx.ID != newEx.ID ) {
-      map.update( newEx.ID, originEx.ID )
-      inheritUID( newEx, originEx.ID )
+  override def onDeclTransformation(
+      originalDecl: TlaDecl,
+      newDecl: TlaDecl
+  ): Unit = {
+    if (originalDecl.ID != newDecl.ID) {
+      map.update(newDecl.ID, originalDecl.ID)
     }
   }
 
   /**
     * Is identifier p_id registered as a result of translation?
+    *
     * @param p_id expression identifier
     * @return true, if p_id has been registered
     */
   def isDefinedAt(p_id: UID): Boolean = {
     map.isDefinedAt(p_id)
   }
 
-  def traceBack( p_id : UID ) : UID = map.get( p_id ) match {
-    case Some( id ) => traceBack( id )
-    case _ => p_id
-  }
+  def traceBack(p_id: UID): UID =
+    map.get(p_id) match {
+      case Some(id) => traceBack(id)
+      case _        => p_id
+    }
 }

tlair/src/main/scala/at/forsyte/apalache/tla/lir/transformations/TransformationListener.scala

@@ -1,15 +1,16 @@
 package at.forsyte.apalache.tla.lir.transformations
 
-import at.forsyte.apalache.tla.lir.TlaEx
+import at.forsyte.apalache.tla.lir.{TlaDecl, TlaEx}
 
 /**
-  * Many processing methods transform a TLA+ expression into another TLA+ expression. To record these changes,
-  * we have introduced this listener. This is a replacement for SourceDB. The clients who are willing to track changes
-  * should implement their listener and register it with a TransformationFactory.
+  * Many processing methods transform a TLA+ expression into another TLA+ expression. Sometimes, we also
+  * transform declarations. To record these changes, we have introduced this listener. The clients who are willing
+  * to track changes should implement their listener and register it with a TransformationFactory.
   *
   * @author Igor Konnov
   */
 trait TransformationListener {
   def onTransformation(originalEx: TlaEx, newEx: TlaEx): Unit
-}
 
+  def onDeclTransformation(originalDecl: TlaDecl, newDecl: TlaDecl): Unit
+}

tlair/src/main/scala/at/forsyte/apalache/tla/lir/transformations/TransformationTracker.scala

@@ -1,6 +1,6 @@
 package at.forsyte.apalache.tla.lir.transformations
 
-import at.forsyte.apalache.tla.lir.TlaEx
+import at.forsyte.apalache.tla.lir.{TlaEx, TlaOperDecl}
 
 /**
   * An implementation of this trait wraps an expression transformer with the code that somehow tracks the transformation.
@@ -9,12 +9,31 @@ import at.forsyte.apalache.tla.lir.TlaEx
   * @author Igor Konnov
   */
 trait TransformationTracker {
+
   /**
     * Decorate a transformation with a tracker.
+    *
     * @param tr an expression transformation
     * @return a new transformation that applies tr and tracks the changes
     */
-  def trackEx( tr: TlaExTransformation): TlaExTransformation
+  def trackEx(tr: TlaExTransformation): TlaExTransformation
+
+  /**
+    * Decorate a declaration transformation with a tracker.
+    *
+    * @param tr a declaration transformation
+    * @return a new declaration transformation that applies tr and tracks the changes
+    */
+  def trackDecl(tr: TlaDeclTransformation): TlaDeclTransformation
+
+  /**
+    * A specialized version of trackDecl, which is tuned to TlaOperDecl. Most transformation apply only to
+    * user-defined operators. That is why we have added a specialized method for it.
+    *
+    * @param tr a declaration transformation
+    * @return a new declaration transformation that applies tr and tracks the changes
+    */
+  def trackOperDecl(tr: TlaOperDecl => TlaOperDecl): TlaOperDecl => TlaOperDecl
 
   /**
     * Sometimes, one has to track a change in a temporary expression that will get transformed into something else
@@ -27,7 +46,31 @@ trait TransformationTracker {
     * @return the new expression
     */
   def hold(from: TlaEx, to: TlaEx): TlaEx = {
-    def tr(f : TlaEx): TlaEx = to
+    def tr(f: TlaEx): TlaEx = to
     trackEx(tr)(from)
   }
+
+  /**
+    * <p>This adapter method takes an expression transformation and turns it into a declaration transformation that:</p>
+    * <ol>
+    *   <li>Copies the declaration body and applies the expression transformation to it, and</li>
+    *   <li>Tracks the update of the declaration.</li>
+    *
+    * <p>We provide this method for convenience, because in many cases the operator body is the only thing that changes
+    *    in an operator declaration.</p>
+    *
+    * @param transformation expression transformation
+    * @return a declaration transformation that is tracking updates to declarations and expressions inside them.
+    */
+  def fromExToDeclTransformation(
+      transformation: TlaExTransformation
+  ): TlaDeclTransformation =
+    trackDecl {
+      case d @ TlaOperDecl(_, _, b) =>
+        d.copy(body = transformation(b))
+
+      case d =>
+        d
+    }
+
 }

tlair/src/main/scala/at/forsyte/apalache/tla/lir/transformations/impl/IdleTracker.scala

@@ -1,6 +1,7 @@
 package at.forsyte.apalache.tla.lir.transformations.impl
 
-import at.forsyte.apalache.tla.lir.transformations.{TlaExTransformation, TransformationTracker}
+import at.forsyte.apalache.tla.lir.TlaOperDecl
+import at.forsyte.apalache.tla.lir.transformations.{TlaDeclTransformation, TlaExTransformation, TransformationTracker}
 
 /**
   * An implementation of TransformationTracker that does not do anything.
@@ -10,4 +11,24 @@ class IdleTracker extends TransformationTracker {
   override def trackEx( transformation: TlaExTransformation): TlaExTransformation = {
     transformation
   }
+
+  /**
+    * Decorate a declaration transformation with a tracker.
+    *
+    * @param transformation a declaration transformation
+    * @return a new declaration transformation that applies tr and tracks the changes
+    */
+  override def trackDecl(transformation: TlaDeclTransformation): TlaDeclTransformation = {
+    transformation
+  }
+
+  /**
+    * A specialized version of trackDecl, which is tuned to TlaOperDecl.
+    *
+    * @param transformation a declaration transformation
+    * @return a new declaration transformation that applies tr and tracks the changes
+    */
+  override def trackOperDecl(transformation: TlaOperDecl => TlaOperDecl): TlaOperDecl => TlaOperDecl = {
+    transformation
+  }
 }