Merge pull request #1357 from hylo-lang/fix-capture-detection

Fix capture detection

Sources/Core/AST/AST+Walk.swift

@@ -236,16 +236,6 @@ extension AST {
     }
   }
 
-  /// Visits `b` and its children in pre-order, notifying `o` when a node is entered or left.
-  public func walk<O: ASTWalkObserver>(functionBody b: FunctionBody, notifying o: inout O) {
-    switch b {
-    case .expr(let e):
-      walk(e, notifying: &o)
-    case .block(let s):
-      walk(s, notifying: &o)
-    }
-  }
-
   // MARK: Declarations
 
   /// Visits the children of `n` in pre-order, notifying `o` when a node is entered or left.
@@ -301,7 +291,7 @@ extension AST {
     walk(roots: n.parameters, notifying: &o)
     walk(n.receiver, notifying: &o)
     walk(n.output, notifying: &o)
-    n.body.map({ walk(functionBody: $0, notifying: &o) })
+    n.body.map({ (b) in walk(b.base, notifying: &o) })
   }
 
   /// Visits the children of `n` in pre-order, notifying `o` when a node is entered or left.
@@ -342,7 +332,7 @@ extension AST {
     _ n: MethodImpl, notifying o: inout O
   ) {
     walk(n.receiver, notifying: &o)
-    n.body.map({ walk(functionBody: $0, notifying: &o) })
+    n.body.map({ (b) in walk(b.base, notifying: &o) })
   }
 
   /// Visits the children of `n` in pre-order, notifying `o` when a node is entered or left.
@@ -397,7 +387,7 @@ extension AST {
     _ n: SubscriptImpl, notifying o: inout O
   ) {
     walk(n.receiver, notifying: &o)
-    n.body.map({ walk(functionBody: $0, notifying: &o) })
+    n.body.map({ (b) in walk(b.base, notifying: &o) })
   }
 
   /// Visits the children of `n` in pre-order, notifying `o` when a node is entered or left.

Sources/Core/AST/Decl/CapturingDecl.swift

@@ -4,4 +4,7 @@ public protocol CapturingDecl: Decl, LexicalScope {
   /// The explicit capture declarations of the entity.
   var explicitCaptures: [BindingDecl.ID] { get }
 
+  /// The part of the declaration that may have implicit captures.
+  var sourcesOfImplicitCaptures: [AnyNodeID] { get }
+
 }

Sources/Core/AST/Decl/FunctionDecl.swift

@@ -130,6 +130,9 @@ public struct FunctionDecl: CapturingDecl, ExposableDecl, GenericDecl, GenericSc
     }
   }
 
+  /// The part of the declaration that may have implicit captures.
+  public var sourcesOfImplicitCaptures: [AnyNodeID] { body.map({ (b) in [b.base] }) ?? [] }
+
   public func validateForm(in ast: AST, reportingDiagnosticsTo log: inout DiagnosticSet) {
     if !isInExprContext {
       // Parameter declarations must have a type annotation.

Sources/Core/AST/Decl/SubscriptDecl.swift

@@ -84,6 +84,9 @@ public struct SubscriptDecl: BundleDecl, CapturingDecl, GenericDecl, GenericScop
   /// Returns whether the declaration denotes a static subscript.
   public var isStatic: Bool { memberModifier?.value == .static }
 
+  /// The part of the declaration that may have implicit captures.
+  public var sourcesOfImplicitCaptures: [AnyNodeID] { impls.map(AnyNodeID.init) }
+
   public func validateForm(in ast: AST, reportingDiagnosticsTo log: inout DiagnosticSet) {
     // Parameter declarations must have a type annotation.
     for p in parameters {

Sources/FrontEnd/AST+UseCollection.swift

@@ -7,9 +7,11 @@ extension AST {
   ///
   /// This method collects all name expressions that occurs in `d`, visiting its children in
   /// pre-order. Nested type and extension declarations are not visited.
-  func uses(in d: AnyDeclID) -> [(NameExpr.ID, AccessEffect)] {
+  func uses<T: CapturingDecl>(in d: T.ID) -> [(NameExpr.ID, AccessEffect)] {
     var v = UseVisitor()
-    walk(d, notifying: &v)
+    for s in self[d].sourcesOfImplicitCaptures {
+      walk(s, notifying: &v)
+    }
     return v.uses
   }
 

Sources/FrontEnd/TypeChecking/TypeChecker.swift

@@ -2602,7 +2602,7 @@ struct TypeChecker {
     }
 
     // Look at uses to update conventions where we could only guess `let` from the context.
-    for (n, m) in program.ast.uses(in: AnyDeclID(program[e].decl)) {
+    for (n, m) in program.ast.uses(in: program[e].decl) {
       let candidates = lookup(unqualified: program[n].name.value.stem, in: program[n].scope)
       guard
         let pick = candidates.unique(ParameterDecl.self),
@@ -2680,7 +2680,7 @@ struct TypeChecker {
   }
 
   /// Returns the implicit captures found in the body of `d`.
-  private mutating func implicitCaptures<T: Decl & LexicalScope>(
+  private mutating func implicitCaptures<T: CapturingDecl>(
     of d: T.ID, ignoring explicitCaptures: Set<String>
   ) -> [TupleType.Element] {
     // Only local declarations have captures.
@@ -2690,7 +2690,7 @@ struct TypeChecker {
     }
 
     var captureToStemAndEffect: [AnyDeclID: (stem: String, effect: AccessEffect)] = [:]
-    for (name, mutability) in program.ast.uses(in: AnyDeclID(d)) {
+    for (name, mutability) in program.ast.uses(in: d) {
       guard
         let (stem, pick) = lookupImplicitCapture(name, occurringIn: d),
         !explicitCaptures.contains(stem)

Tests/EndToEndTests/TestCases/ExplicitCaptures.hylo

@@ -0,0 +1,15 @@
+//- compileAndRun expecting: success
+
+fun apply<E>(_ f: [E]() -> Int) -> Int {
+  f()
+}
+
+public fun main() {
+  var local_variable = 0
+  let p = mutable_pointer[to: &local_variable]
+  let n = apply(fun[sink let q = p.copy()]() -> Int {
+    &(q.unsafe[]) = 19
+    return 19
+  })
+  precondition(n == local_variable)
+}