address lambda collision issue

johnynek · Feb 5, 2024 · f6629bc · f6629bc
1 parent d2778e9
commit f6629bc
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 1 deletion.
diff --git a/cli/src/test/scala/org/bykn/bosatsu/codegen/python/CodeTest.scala b/cli/src/test/scala/org/bykn/bosatsu/codegen/python/CodeTest.scala
@@ -475,6 +475,22 @@ else:
     }
   }
 
+  test("(lambda x: lambda y: x + y)(y)") {
+    val x = Code.Ident("x")
+    val y = Code.Ident("y")
+    val hardCase = Code.Lambda(List(x),
+      Code.Lambda(List(y), x + y)
+    )
+
+    val applied = hardCase(y).simplify
+    val y0 = Code.Ident("y0")
+    assert(applied == Code.Lambda(List(y0), y + y0))
+
+    val z = Code.Ident("z")
+    val applied1 = hardCase(z).simplify
+    assert(applied1 == Code.Lambda(List(y), z + y))
+  }
+
   test("simplify(Map.empty, x) == x") {
     forAll(genExpr(4)) { x =>
       assert(Code.substitute(Map.empty, x) == x)  

diff --git a/core/src/main/scala/org/bykn/bosatsu/codegen/python/Code.scala b/core/src/main/scala/org/bykn/bosatsu/codegen/python/Code.scala
@@ -659,13 +659,77 @@ object Code {
         val argsSet = args.toSet
         val sm1 = subMap.filterNot { case (i, _) => argsSet(i) }
         if (sm1.isEmpty) in
-        else Lambda(args, substitute(sm1, result))
+        else {
+          // reduce is safe here because sm1.nonEmpty
+          val subFrees = sm1.iterator.map { case (_, v) => freeIdents(v) }.reduce(_ | _)
+          val clashIdent = argsSet & subFrees
+          if (clashIdent.isEmpty) {
+            Lambda(args, substitute(sm1, result))
+          }
+          else {
+            // we have to allocate new variables
+            def alloc(rename: List[Ident], avoid: Set[Ident]): List[Ident] =
+              rename match {
+                case Nil => Nil
+                case (i @ Ident(nm)) :: tail => 
+                  if (clashIdent(i)) {
+                    val nm1 =
+                      Iterator.from(0).map { i => Ident(nm + i.toString) }
+                        .collectFirst { case n if !avoid(n) => n }
+                        .get
+                    nm1 :: alloc(tail, avoid + nm1)
+                  }
+                  else {
+                    i :: alloc(tail, avoid + i)
+                  }
+              }
+
+            val avoids = subFrees | freeIdents(result) | sm1.keySet
+            val newArgs = alloc(args, avoids)
+            val resSub = args.zip(newArgs).toMap
+            val res1 = substitute(resSub, result)
+            substitute(sm1, Lambda(newArgs, res1))
+          }
+        }
       case Apply(fn, args) =>
         Apply(substitute(subMap, fn), args.map(substitute(subMap, _)))
       case DotSelect(ex, ident) =>
         DotSelect(substitute(subMap, ex), ident)
     }
 
+  def freeIdents(ex: Expression): Set[Ident] = {
+    def loop(ex: Expression, bound: Set[Ident]): Set[Ident] =
+      ex match {
+        case PyInt(_) | PyString(_) | PyBool(_) => Set.empty
+        case i@Ident(_) =>
+          if (bound(i)) Set.empty
+          else Set(i)
+        case Op(left, _, right) => loop(left, bound) | loop(right, bound)
+        case Parens(expr) =>
+          loop(expr, bound)
+        case SelectItem(arg, position) =>
+          loop(arg, bound) | loop(position, bound)
+        case SelectRange(arg, start, end) =>
+          loop(arg, bound) |
+            start.map(loop(_, bound)).getOrElse(Set.empty) |
+            end.map(loop(_, bound)).getOrElse(Set.empty)
+        case Ternary(ifTrue, cond, ifFalse) =>
+          loop(ifTrue, bound) | loop(cond, bound) | loop(ifFalse, bound)
+        case MakeTuple(args) =>
+          args.foldLeft(Set.empty[Ident])(_ | loop(_, bound))
+        case MakeList(args) =>
+          args.foldLeft(Set.empty[Ident])(_ | loop(_, bound))
+        case Lambda(args, result) =>
+          loop(result, bound ++ args)
+        case Apply(fn, args) =>
+          loop(fn, bound) | args.foldLeft(Set.empty[Ident])(_ | loop(_, bound))
+        case DotSelect(ex, _) =>
+          loop(ex, bound)
+      }
+
+    loop(ex, Set.empty)
+  }
+
   def toReturn(v: ValueLike): Statement =
     v match {
       case x: Expression => Code.Return(x)