Generate raw IR for calls to method bundles

Sources/FrontEnd/TypedProgram.swift

@@ -428,7 +428,15 @@ public struct TypedProgram {
     }
   }
 
-  /// Returns the declarations of `d`' captures.
+  /// Returns the type of `d` specialized by `specialization` in `scopeOfUse`.
+  public func canonicalType<T: Decl>(
+    of d: T.ID, specializedBy specialization: GenericArguments, in scopeOfUse: AnyScopeID
+  ) -> AnyType {
+    let t = specialize(self[d].type, for: specialization, in: scopeOfUse)
+    return canonical(t, in: scopeOfUse)
+  }
+
+  /// Returns the declarations of `d`'s captures.
   ///
   /// If `d` is a member, its receiver is its only capture. Otherwise, this method returns
   /// `nonMemberCaptures(d)`.
@@ -439,7 +447,7 @@ public struct TypedProgram {
     return nonMemberCaptures(of: d)
   }
 
-  /// Returns the declarations of `d`' captures.
+  /// Returns the declarations of `d`'s captures.
   ///
   /// If `d` is a member, its receiver is its only capture. Otherwise, this method returns
   /// `nonMemberCaptures(d)`.

Sources/IR/Analysis/Module+AccessReification.swift

@@ -145,7 +145,7 @@ extension Module {
       arguments[a] = .register(insert(b, before: i))
     }
 
-    let o = RemoteType(k, s.projection.bareType)
+    let o = RemoteType(k, s.projection)
     let reified = makeProject(
       o, applying: s.variants[k]!, specializedBy: s.bundle.arguments, to: arguments, at: s.site)
     replace(i, with: reified)

Sources/IR/BundleReference.swift

@@ -0,0 +1,30 @@
+import Core
+
+/// A reference to a method or subscript bundle.
+public struct BundleReference<T: BundleDecl>: Hashable {
+
+  /// The ID of the referred bundle.
+  public let bundle: T.ID
+
+  /// If `bundle` is generic, the arguments to its generic parameter.
+  public let arguments: GenericArguments
+
+  /// Creates a reference to `s` parameterized by `a`.
+  public init(to s: T.ID, specializedBy a: GenericArguments) {
+    self.bundle = s
+    self.arguments = a
+  }
+
+}
+
+extension BundleReference: CustomStringConvertible {
+
+  public var description: String {
+    if arguments.isEmpty {
+      return "@\(bundle)"
+    } else {
+      return "@\(bundle)<\(list: arguments.values)>"
+    }
+  }
+
+}

Sources/IR/Callee.swift

@@ -0,0 +1,30 @@
+import Core
+
+/// The callee of a function call in the IR.
+enum Callee {
+
+  /// A direct reference to a function, initializer, or method implementation.
+  case direct(FunctionReference)
+
+  /// A reference to a variant in a method bundle.
+  case bundle(BundleReference<MethodDecl>)
+
+  /// A lambda.
+  case lambda(Operand)
+
+  /// Creates an instance representing a reference to the lowered form of `d` in `module`,
+  /// specialized by`a` in `scopeOfUse`.
+  init(
+    _ d: AnyDeclID, specializedBy a: GenericArguments,
+    in module: inout Module, usedIn scopeOfUse: AnyScopeID
+  ) {
+    if let m = MethodDecl.ID(d) {
+      self = .bundle(BundleReference(to: m, specializedBy: a))
+    } else {
+      let r = FunctionReference(
+        to: FunctionDecl.ID(d)!, in: &module, specializedBy: a, in: scopeOfUse)
+      self = .direct(r)
+    }
+  }
+
+}

Sources/IR/Emitter.swift

@@ -1164,8 +1164,14 @@ struct Emitter {
     let arguments = captures + explicitArguments
 
     // Call is evaluated last.
-    let o = insert(module.makeAccess(.set, from: storage, at: ast[e].site))!
-    insert(module.makeCall(applying: callee, to: arguments, writingResultTo: o, at: ast[e].site))
+    switch callee {
+    case .direct(let r):
+      emitApply(.constant(r), to: arguments, writingResultTo: storage, at: ast[e].site)
+    case .lambda(let r):
+      emitApply(r, to: arguments, writingResultTo: storage, at: ast[e].site)
+    case .bundle(let r):
+      emitApply(r, to: arguments, writingResultTo: storage, at: ast[e].site)
+    }
   }
 
   /// Inserts the IR for storing the value of `e` to `storage`.
@@ -1404,6 +1410,27 @@ struct Emitter {
     frames.top.setMayHoldCaptures(s)
   }
 
+  /// Inserts the IR for calling `callee` on `arguments`, storing the result to `storage`.
+  private mutating func emitApply(
+    _ callee: Operand, to arguments: [Operand],
+    writingResultTo storage: Operand, at site: SourceRange
+  ) {
+    let o = insert(module.makeAccess(.set, from: storage, at: site))!
+    insert(module.makeCall(applying: callee, to: arguments, writingResultTo: o, at: site))
+    insert(module.makeEndAccess(o, at: site))
+  }
+
+  /// Inserts the IR for calling `callee` on `arguments`, storing the result to `storage`.
+  private mutating func emitApply(
+    _ callee: BundleReference<MethodDecl>, to arguments: [Operand],
+    writingResultTo storage: Operand, at site: SourceRange
+  ) {
+    let o = insert(module.makeAccess(.set, from: storage, at: site))!
+    insert(module.makeCallBundle(
+      applying: .init(callee, in: insertionScope!), to: arguments, writingResultTo: o, at: site))
+    insert(module.makeEndAccess(o, at: site))
+  }
+
   /// Inserts the IR for given constructor `call`, which initializes storage `r` by applying
   /// initializer `d` parameterized by `a`.
   ///
@@ -1579,7 +1606,7 @@ struct Emitter {
   /// - Requires: `callee` has a lambda type.
   private mutating func emit(
     functionCallee callee: AnyExprID
-  ) -> (callee: Operand, captures: [Operand]) {
+  ) -> (callee: Callee, captures: [Operand]) {
     switch callee.kind {
     case NameExpr.self:
       return emit(namedFunctionCallee: .init(callee)!)
@@ -1589,41 +1616,36 @@ struct Emitter {
       return emit(functionCallee: ast[InoutExpr.ID(callee)!].subject)
 
     default:
-      return (emit(lambdaCallee: callee), [])
+      let f = emit(lambdaCallee: callee)
+      return (.lambda(f), [])
     }
   }
 
   /// Inserts the IR for given `callee` and returns `(c, a)`, where `c` is the callee's value and
   /// `a` are arguments to lifted parameters.
-  ///
-  /// - Requires: `callee` has a lambda type.
   private mutating func emit(
     namedFunctionCallee callee: NameExpr.ID
-  ) -> (callee: Operand, captures: [Operand]) {
-    let calleeType = LambdaType(canonical(program[callee].type))!
-
+  ) -> (callee: Callee, captures: [Operand]) {
     switch program[callee].referredDecl {
     case .direct(let d, let a) where d.kind == FunctionDecl.self:
       // Callee is a direct reference to a function declaration.
-      guard calleeType.environment == .void else {
+      guard LambdaType(canonical(program[callee].type))!.environment == .void else {
         UNIMPLEMENTED()
       }
 
       let specialization = module.specialization(in: insertionFunction!).merging(a)
-      let r = FunctionReference(
+      let f = FunctionReference(
         to: FunctionDecl.ID(d)!, in: &module, specializedBy: specialization, in: insertionScope!)
-      return (.constant(r), [])
+      return (.direct(f), [])
 
-    case .member(let d, let a, let s) where d.kind == FunctionDecl.self:
-      // Callee is a member reference to a function or method.
-      let r = FunctionReference(
-        to: FunctionDecl.ID(d)!, in: &module, specializedBy: a, in: insertionScope!)
-
-      // The callee's receiver is the sole capture.
+    case .member(let d, let a, let s):
+      // Callee is a member reference to a function or method. Its receiver is the only capture.
+      let specialization = module.specialization(in: insertionFunction!).merging(a)
       let receiver = emitLValue(receiver: s, at: ast[callee].site)
-      let k = RemoteType(calleeType.captures[0].type)?.access ?? .sink
-      let i = insert(module.makeAccess(k, from: receiver, at: ast[callee].site))!
-      return (Operand.constant(r), [i])
+      let k = receiverCapabilities(program[callee].type)
+      let c = insert(module.makeAccess(k, from: receiver, at: ast[callee].site))!
+      let f = Callee(d, specializedBy: specialization, in: &module, usedIn: insertionScope!)
+      return (f, [c])
 
     case .builtinFunction, .builtinType:
       // Calls to built-ins should have been handled already.
@@ -1632,7 +1654,7 @@ struct Emitter {
     default:
       // Callee is a lambda.
       let f = emit(lambdaCallee: .init(callee))
-      return (f, [])
+      return (.lambda(f), [])
     }
   }
 
@@ -1655,7 +1677,7 @@ struct Emitter {
 
   private mutating func emit(
     subscriptCallee callee: AnyExprID
-  ) -> (callee: SubscriptBundleReference, captures: [Operand]) {
+  ) -> (callee: BundleReference<SubscriptDecl>, captures: [Operand]) {
     // TODO: Handle captures
     switch callee.kind {
     case NameExpr.self:
@@ -1672,7 +1694,7 @@ struct Emitter {
 
   private mutating func emit(
     namedSubscriptCallee callee: NameExpr.ID
-  ) -> (callee: SubscriptBundleReference, captures: [Operand]) {
+  ) -> (callee: BundleReference<SubscriptDecl>, captures: [Operand]) {
     switch program[callee].referredDecl {
     case .direct(let d, let a) where d.kind == SubscriptDecl.self:
       // Callee is a direct reference to a subscript declaration.
@@ -1681,12 +1703,12 @@ struct Emitter {
         UNIMPLEMENTED()
       }
 
-      let b = SubscriptBundleReference(to: SubscriptDecl.ID(d)!, parameterizedBy: a)
+      let b = BundleReference(to: SubscriptDecl.ID(d)!, specializedBy: a)
       return (b, [])
 
     case .member(let d, let a, let s) where d.kind == SubscriptDecl.self:
       // Callee is a member reference to a subscript declaration.
-      let b = SubscriptBundleReference(to: SubscriptDecl.ID(d)!, parameterizedBy: a)
+      let b = BundleReference(to: SubscriptDecl.ID(d)!, specializedBy: a)
 
       // The callee's receiver is the sole capture.
       let receiver = emitLValue(receiver: s, at: ast[callee].site)
@@ -2034,17 +2056,9 @@ struct Emitter {
     let (callee, captures) = emit(subscriptCallee: ast[e].callee)
     let arguments = captures + explicitArguments
 
-    var variants: [AccessEffect: Function.ID] = [:]
-    for v in ast[callee.bundle].impls {
-      variants[ast[v].introducer.value] = module.demandDeclaration(lowering: v)
-    }
-
-    // Projection is evaluated last.
-    let t = canonicalType(of: callee.bundle, specializedBy: callee.arguments)
-    return insert(
-      module.makeProjectBundle(
-        applying: variants, of: callee, typed: SubscriptType(t)!,
-        to: arguments, at: ast[e].site))!
+    let s = module.makeProjectBundle(
+      applying: .init(callee, in: insertionScope!), to: arguments, at: ast[e].site)
+    return insert(s)!
   }
 
   /// Inserts the IR for lvalue `e`.
@@ -2145,17 +2159,12 @@ struct Emitter {
         boundTo: receiver, declaredBy: i, specializedBy: specialization, at: site)
     }
 
+    let callee = BundleReference(to: d, specializedBy: specialization)
     let t = SubscriptType(canonicalType(of: d, specializedBy: specialization))!
     let r = insert(module.makeAccess(t.capabilities, from: receiver, at: site))!
 
-    var variants: [AccessEffect: Function.ID] = [:]
-    for v in ast[d].impls {
-      variants[ast[v].introducer.value] = module.demandDeclaration(lowering: v)
-    }
-
     let s = module.makeProjectBundle(
-      applying: variants, of: .init(to: d, parameterizedBy: specialization), typed: t,
-      to: [r], at: site)
+      applying: .init(callee, in: insertionScope!), to: [r], at: site)
     return insert(s)!
   }
 
@@ -2407,6 +2416,19 @@ struct Emitter {
     return canonical(t)
   }
 
+  /// Returns the capabilities required by the receiver of `callee`, which is the type of a member
+  /// function or method bundle.
+  private func receiverCapabilities(_ callee: AnyType) -> AccessEffectSet {
+    switch canonical(callee).base {
+    case let t as LambdaType:
+      return [RemoteType(t.captures[0].type)?.access ?? .sink]
+    case let t as MethodType:
+      return t.capabilities
+    default:
+      unreachable()
+    }
+  }
+
   /// Inserts a stack allocation for an object of type `t`.
   private mutating func emitAllocStack(
     for t: AnyType, at site: SourceRange

Sources/IR/FunctionID.swift

@@ -1,4 +1,5 @@
 import Core
+import Utils
 
 extension Function {
 
@@ -28,23 +29,21 @@ extension Function {
     /// The value of this identity.
     public let value: Value
 
-    /// Creates the identity of the lowered form of `f`.
-    public init(_ f: FunctionDecl.ID) {
-      self.value = .lowered(AnyDeclID(f))
-    }
-
-    /// Creates the identity of the lowered form of `f` used as an initializer.
-    public init(initializer f: InitializerDecl.ID) {
-      self.value = .lowered(AnyDeclID(f))
-    }
-
-    /// Creates the identity of the lowered form of `s`.
-    public init(_ s: SubscriptImpl.ID) {
-      self.value = .loweredSubscript(s)
+    /// Creates the identity of the lowered form of `d`, which is the declaration of a function,
+    /// initializer, method implementation, or subscript implementation.s
+    init<T: DeclID>(_ d: T) {
+      switch d.kind {
+      case FunctionDecl.self, InitializerDecl.self, MethodImpl.self:
+        self.value = .lowered(AnyDeclID(d))
+      case SubscriptImpl.self:
+        self.value = .loweredSubscript(SubscriptImpl.ID(d)!)
+      default:
+        unreachable()
+      }
     }
 
     /// Creates the identity of the lowered form of `s`.
-    public init(_ s: SynthesizedFunctionDecl) {
+    init(_ s: SynthesizedFunctionDecl) {
       self.value = .synthesized(s)
     }