Extract statement-to-expression from PureMethodsToFunctions

src/col/vct/col/util/StatementToExpression.scala

@@ -0,0 +1,113 @@
+package vct.col.util
+
+import vct.col.ast.{
+  Assign,
+  Block,
+  BooleanValue,
+  Branch,
+  Expr,
+  Let,
+  Local,
+  Return,
+  Scope,
+  Select,
+  Statement,
+  Variable,
+}
+import vct.col.origin.{DiagnosticOrigin, Origin}
+import vct.col.rewrite.{Generation, NonLatchingRewriter}
+import vct.result.VerificationError
+
+case object StatementToExpression {
+
+  def toExpression[Pre <: Generation, Post <: Generation](
+      rw: NonLatchingRewriter[Pre, Post],
+      errorBuilder: String => VerificationError,
+      stat: Statement[Pre],
+      alt: => Option[Expr[Post]],
+  ): Option[Expr[Post]] = {
+    implicit val o: Origin = DiagnosticOrigin
+    stat match {
+      case Return(e) =>
+        alt match {
+          case Some(_) =>
+            throw errorBuilder("Dead code after return is not allowed.")
+          case None => Some(rw.dispatch(e))
+        }
+      case Block(Nil) => alt
+      case Block(stat +: tail) =>
+        toExpression(
+          rw,
+          errorBuilder,
+          stat,
+          toExpression(rw, errorBuilder, Block(tail), alt),
+        )
+      case Branch(Nil) => alt
+      case Branch((BooleanValue(true), impl) +: _) =>
+        toExpression(rw, errorBuilder, impl, alt)
+      case Branch((cond, impl) +: branches) =>
+        Some(Select(
+          rw.dispatch(cond),
+          toExpression(rw, errorBuilder, impl, alt).getOrElse(return None),
+          toExpression(rw, errorBuilder, Branch(branches), alt)
+            .getOrElse(return None),
+        ))
+      case Scope(locals, impl) =>
+        if (!locals.forall(countAssignments(_, impl).exists(_ <= 1))) {
+          return None
+        }
+        toExpression(rw, errorBuilder, impl, alt)
+      case Assign(Local(ref), e) =>
+        rw.localHeapVariables.scope {
+          rw.variables.scope {
+            alt match {
+              case Some(exprAlt) =>
+                Some(Let[Post](
+                  rw.variables.collect { rw.dispatch(ref.decl) }._1.head,
+                  rw.dispatch(e),
+                  exprAlt,
+                ))
+              case None =>
+                throw errorBuilder("Assign may not be the last statement")
+            }
+          }
+        }
+
+      case _ => None
+    }
+  }
+
+  def countAssignments[G <: Generation](
+      v: Variable[G],
+      s: Statement[G],
+  ): Option[Int] =
+    s match {
+      case Return(_) => Some(0)
+      case Block(stats) =>
+        val x =
+          stats.map(countAssignments(v, _)).fold(Some(0)) {
+            case (Some(a), Some(b)) => Some(a + b)
+            case _ => None
+          }
+        x
+      case Branch(conds) =>
+        val assignmentCounts = conds.map(_._2).map(countAssignments(v, _))
+          .collect {
+            case Some(n) => n
+            case None => return None
+          }
+        if (assignmentCounts.forall(_ <= 1)) {
+          assignmentCounts.maxOption.orElse(Some(0))
+        } else { None }
+      case Assign(Local(ref), _) =>
+        Some(
+          if (ref.decl == v)
+            1
+          else
+            0
+        )
+      case Assign(_, _) => None
+      case _ => None
+    }
+
+}

src/rewrite/vct/rewrite/PureMethodsToFunctions.scala

@@ -2,12 +2,9 @@ package vct.col.rewrite
 
 import hre.util.ScopedStack
 import vct.col.ast._
-import vct.col.origin.{DiagnosticOrigin, LabelContext, Origin, PreferredName}
-import vct.col.rewrite.{Generation, Rewriter, RewriterBuilder}
-import vct.col.ast.RewriteHelpers._
+import vct.col.origin.{DiagnosticOrigin, LabelContext, Origin}
 import vct.col.ref.Ref
-import vct.col.util.AstBuildHelpers._
-import vct.col.util.SuccessionMap
+import vct.col.util.StatementToExpression
 import vct.result.VerificationError.UserError
 
 case object PureMethodsToFunctions extends RewriterBuilder {
@@ -33,6 +30,7 @@ case class PureMethodsToFunctions[Pre <: Generation]() extends Rewriter[Pre] {
 
   val currentAbstractMethod: ScopedStack[AbstractMethod[Pre]] = ScopedStack()
 
+  /*
   def countAssignments(v: Variable[Pre], s: Statement[Pre]): Option[Int] =
     s match {
       case Return(_) => Some(0)
@@ -62,7 +60,9 @@ case class PureMethodsToFunctions[Pre <: Generation]() extends Rewriter[Pre] {
       case Assign(_, _) => None
       case _ => None
     }
+   */
 
+  /*
   def toExpression(
       stat: Statement[Pre],
       alt: => Option[Expr[Post]],
@@ -116,6 +116,7 @@ case class PureMethodsToFunctions[Pre <: Generation]() extends Rewriter[Pre] {
       case _ => None
     }
   }
+   */
 
   override def dispatch(decl: Declaration[Pre]): Unit = {
     implicit val o: Origin = decl.o
@@ -137,7 +138,12 @@ case class PureMethodsToFunctions[Pre <: Generation]() extends Rewriter[Pre] {
             body =
               currentAbstractMethod.having(proc) {
                 proc.body.map(
-                  toExpression(_, None).getOrElse(
+                  StatementToExpression.toExpression(
+                    this,
+                    (s: String) => MethodCannotIntoFunction(proc, s),
+                    _,
+                    None,
+                  ).getOrElse(
                     throw MethodCannotIntoFunction(
                       proc,
                       "the method implementation cannot be restructured into a pure expression",
@@ -169,7 +175,12 @@ case class PureMethodsToFunctions[Pre <: Generation]() extends Rewriter[Pre] {
             body =
               currentAbstractMethod.having(method) {
                 method.body.map(
-                  toExpression(_, None).getOrElse(
+                  StatementToExpression.toExpression(
+                    this,
+                    (s: String) => MethodCannotIntoFunction(method, s),
+                    _,
+                    None,
+                  ).getOrElse(
                     throw MethodCannotIntoFunction(
                       method,
                       "the method implementation cannot be restructured into a pure expression",