ML Backend yields different results than JS backend after monomorphization

[hflsmax]

I'm trying to implement a triples program (pasted at the end) that's similar to this one https://github.com/koka-lang/libmprompt/blob/main/test/src/triples.c

This program can be compiled to ml without error, but mlton can't compiled the generated ml code. It gives the following error. Did I make a mistake in my source program or is it a bug in the Effekt compiler?

mlton main.mlb
Error: triples.sml 1180.11-1230.32.
  Function applied to incorrect argument.
    expects: _ -> ([unit] -> _) -> _
    but got: _ -> ([int] -> _) -> _
    in: ((fn ev1409_1409 => (fn ChoiceDol  ...   => (k21682 a1681)))
Error: triples.sml 1242.14-1265.28.
  Function applied to incorrect argument.
    expects: (_ -> _ -> _ -> _ -> ([int] -> _) -> _) Yield_759
    but got: (_ -> _ -> _ -> _ -> ([unit] -> _) -> _) Yield_759
    in: ((fn ev1414_1414 => (fn YieldDoll  ...  d end)) k1686)))))))  

import io/args
import immutable/list
import immutable/option
import text/string

effect Choice {
  def choose (n: Int): Int
  def fail(): Unit
}

effect Yield {
  def yield (a: Int, b: Int, c: Int): Unit
}

def triple(n: Int, s: Int): Unit / {Choice, Yield} = {
  val i = do choose(n)
  val j = do choose(i - 1)
  val k = do choose(j - 1)
  if (i + j + k == s) {
    do yield(i, j, k)
  } else {
    do fail()
  }
}

def yieldTriples(n: Int, s: Int): Unit / {Yield} = {
  try {
    triple(n, s)
  } with Choice {
    def choose(n) = {
      var i = 1;
      while (i != n) {
        resume(i);
        i = i + 1
      }
    }
    def fail() = ()
  }
}

def countTriples(n: Int, s: Int): Int = {
  var cnt = 0;
  try {
    yieldTriples(n, s)
  } with Yield {
    def yield(i: Int, j: Int, k: Int) = {
      cnt = cnt + 1; 
      resume(())
    }
  }; cnt
}

def main() = println(countTriples(8, 6))

[b-studios]

I am sorry, but with the MLton backend you are working on the bleeding edge. This is indeed a known issue that we still need to address. ML does not have support for rank-2 types, which are required by our translation to system-f. We are working on #229 to avoid this.

~The easiest workaround is avoiding the interaction of mutable state, recursive functions (while) and calling the continuation under the recursive function for now. This is not really satisfying and we are working on better solutions.~

EDIT: I believe the issue was actually not caused by mutable state and recursive functions, but by the fact that Choice contained both effect operations choose and fail. They are naturally called at different answer types (int and unit, the result of the operations). MLton will try to monomorphize the evidence passed for Choice but it cannot, because it would need to work for int as well as unit (this problem is dealt with in #229 -- but some other problems resurfaced there).

I slightly massaged your program manually to avoid the rank-2 problem and it should work now:

effect Choose {
  def choose(n: Int): Int
}
effect Fail {
  def fail(): Unit
}
effect Choice = {Choose,Fail}

effect Yield {
  def yield(a: Int, b: Int, c: Int): Unit
}

def triple(n: Int, s: Int): Unit / {Choice, Yield} = {
  val i = do choose(n)
  val j = do choose(i - 1)
  val k = do choose(j - 1)
  if (i + j + k == s) {
    do yield(i, j, k)
  } else {
    do fail()
  }
}

def yieldTriples(n: Int, s: Int): Unit / {Yield} = {
  try {
    triple(n, s)
  } with Choose {
    def choose(n) = {
      var i = 1;
      while (i < n) {
        resume(i);
        i = i + 1
      }
    }
  } with Fail {
    def fail() = ()
  }
}

def countTriples(n: Int, s: Int): Int = {
  var cnt = 0;
  try {
    yieldTriples(n, s)
  } with Yield {
    def yield(i: Int, j: Int, k: Int) = {
      cnt = cnt + 1;
      resume(())
    }
  };
  cnt
}

def main() = println(countTriples(8, 6))

[b-studios]

BTW, if you are benchmarking this; I would be interested to learn about your findings :)

[b-studios]

Now that #229 is merged, we should revisit this issue. Actually, I just tried it and the ml implementation gives 0 while js gives 4 for println(countTriples(10, 14)), which needs to be a bug.

[mm0821]

I guess that this is because lift inference is not properly implemented for continuations. In order to call continuations under further handlers, they should receive evidence just as capabilities do, but in the implementation they do not. Here is smaller example:

effect Foo(): Int
effect Foo2(x: Int): Int
effect Foo3(x: Int): Int

def main() = println(try {
  val z = try {
    val y = do Foo();
    do Foo2(y)
  } with Foo { () =>
    try {
      resume(2)
    } with Foo3 { (x: Int) =>
      42
    }
  }
  z + 3
} with Foo2 { (x: Int) =>
  resume(x) + resume(x) + resume(x)
})

The result of this should be (2+3) + (2+3) + (2+3) = 15 which is the case in the JS backend. But the mlton backend yields 9. This is because the continuation for Foo is called under a further handler, and hence so is the invocation of Foo2, since the continuation of Foo is not lifted. Thus, the continuation captured for Foo2 does not include the addition of 3 which means that 3 is added only once in the end, resulting in 2 + 2 + 2 + 3 = 9.

Something similar happens in the Triples example in yieldTriples. There the continuation in the Choose handler is not directly called under a further handler, but in the context of local mutable state which is translated to a handler for a state effect. I suppose this is the reason why the Triples example goes wrong in the mlton backend. But I have not fully thought through it, so there might also be another cause. Still, I wanted to report this.