From 5178f0a26be61f0a6dc81d5e68a5014349554261 Mon Sep 17 00:00:00 2001
From: zerbina <100542850+zerbina@users.noreply.github.com>
Date: Wed, 21 Aug 2024 20:03:12 +0200
Subject: [PATCH] sempass2: lift capture restriction (#1427)

## Summary

Procedures useable in both compile- and run-time contexts can now close
over compile-time-only locals even if nested in another procedure
usable in both contexts.

Fixes https://github.com/nim-works/nimskull/issues/1425.

## Details

Given the nested procedures `a: b: c:`, where `c` closes over a local
of `a`, `c` effectively captures from `b` and `b` captures from `a`.

A procedure usable in both a compile- and run-time context may capture
from a compile-time-only procedure, so the aforementioned should work
when `a` is a compile-time-only procedure and the others are not, since
`c` can legally borrow from `b` and `b` from `a`, but it was disallowed
in practice.

The arbitrary restriction is now lifted, and the code implementing the
test restructured to be easier to understand.
---
 compiler/sem/sempass2.nim                | 26 ++++++++++++++----------
 tests/lang_callable/closure/tclosure.nim | 18 ++++++++++++++++
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/compiler/sem/sempass2.nim b/compiler/sem/sempass2.nim
index 306bb52febd..291f43d34e4 100644
--- a/compiler/sem/sempass2.nim
+++ b/compiler/sem/sempass2.nim
@@ -1684,22 +1684,26 @@ proc canCaptureFrom*(captor, target: PSym): bool =
   ## taking the compile-time/run-time boundary into account:
   ## 1) attempting to capture a local defined outside an inner macro from
   ##   inside the macro is illegal
-  ## 2) closing over a local defined inside a compile-time-only routine from a
-  ##   routine than can also be used at run-time is only valid if the chain of
-  ##   enclosing routines leading up to `target` are all compile-time-only
+  ## 2) routines usable at run-time may close over locals of compile-time-only
+  ##   routines
   ## 3) a compile-time-only routine closing over a run-time location is illegal
   template isCompTimeOnly(s: PSym): bool =
     sfCompileTime in s.flags
 
-  result = not captor.isCompTimeOnly or target.isCompTimeOnly    # rule #3
-  # check `captor` and all enclosing routines up to, but not including,
-  # `target` for rule violations
   var s = captor
-  while result and s != target:
-    result =
-      s.kind != skMacro and                                      # rule #1
-      (s == captor or s.isCompTimeOnly == target.isCompTimeOnly) # rule #2
-    s = s.skipGenericOwner
+  while s != target:
+    if s.kind == skMacro:
+      break # rule #1
+
+    let parent = s.skipGenericOwner
+    # going from 'everywhere' to 'compile-time' is disallowed, everything else
+    # is fine
+    if s.isCompTimeOnly and not parent.isCompTimeOnly:
+      break # rule #2 and #3
+
+    s = parent
+
+  result = s == target
 
 # ------------- public interface ----------------
 
diff --git a/tests/lang_callable/closure/tclosure.nim b/tests/lang_callable/closure/tclosure.nim
index 005de5a0037..5c12759c05b 100644
--- a/tests/lang_callable/closure/tclosure.nim
+++ b/tests/lang_callable/closure/tclosure.nim
@@ -457,6 +457,24 @@ block close_over_compile_time_loc:
   static:
     p()
 
+block close_over_compile_time_loc_2:
+  # nested non-compile-time-only procedures can close over locals of compile-
+  # time-only procedures
+  proc p() {.compileTime.} =
+    var x = 0
+    proc inner(cmp: int) = # `inner` is explicitly not compile-time-only
+      proc innerInner(cmp: int) =
+        inc x
+        doAssert x == cmp
+
+      innerInner(cmp)
+
+    inner(1)
+    inner(2)
+
+  static:
+    p()
+
 template test(body: untyped) {.dirty.} =
   ## Tests that `body` works when placed in:
   ## - a normal procedure