Generate IR for calls to method bundles

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), [])
     }
   }
 
@@ -2394,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

@@ -30,10 +30,10 @@ extension Function {
     public let value: Value
 
     /// Creates the identity of the lowered form of `d`, which is the declaration of a function,
-    /// initializer, or subscript implementation.s
+    /// initializer, method implementation, or subscript implementation.s
     init<T: DeclID>(_ d: T) {
       switch d.kind {
-      case FunctionDecl.self, InitializerDecl.self:
+      case FunctionDecl.self, InitializerDecl.self, MethodImpl.self:
         self.value = .lowered(AnyDeclID(d))
       case SubscriptImpl.self:
         self.value = .loweredSubscript(SubscriptImpl.ID(d)!)

Sources/IR/Module.swift

@@ -280,6 +280,30 @@ public struct Module {
     return f
   }
 
+  /// Returns the identity of the IR function corresponding to `d`.
+  mutating func demandDeclaration(lowering d: MethodImpl.ID) -> Function.ID {
+    let f = Function.ID(d)
+    if functions[f] != nil { return f }
+
+    let parameters = program.accumulatedGenericParameters(in: d)
+    let output = program.canonical(
+      (program[d].type.base as! CallableType).output, in: program[d].scope)
+
+    let inputs = loweredParameters(of: d)
+
+    let entity = Function(
+      isSubscript: false,
+      site: program.ast[d].site,
+      linkage: program.isExported(d) ? .external : .module,
+      genericParameters: Array(parameters),
+      inputs: inputs,
+      output: output,
+      blocks: [])
+    addFunction(entity, for: f)
+
+    return f
+  }
+
   /// Returns the identity of the IR function corresponding to `d`.
   mutating func demandDeclaration(lowering d: SubscriptImpl.ID) -> Function.ID {
     let f = Function.ID(d)
@@ -407,6 +431,14 @@ public struct Module {
     return result
   }
 
+  /// Returns the lowered declarations of `d`'s parameters.
+  private func loweredParameters(of d: MethodImpl.ID) -> [Parameter] {
+    let bundle = MethodDecl.ID(program[d].scope)!
+    let inputs = program.ast[bundle].parameters
+    let r = Parameter(AnyDeclID(program[d].receiver), capturedAs: program[d].receiver.type)
+    return [r] + inputs.map(pairedWithLoweredType(parameter:))
+  }
+
   /// Returns the lowered declarations of `d`'s parameters.
   ///
   /// `d`'s receiver comes first and is followed by `d`'s formal parameters, from left to right.

Sources/IR/Operands/Instruction/CallBundle.swift

@@ -0,0 +1,91 @@
+import Core
+import Utils
+
+/// Invokes one variant of `bundle` with `arguments` and writes its result to `output`.
+public struct CallBundle: Instruction {
+
+  /// The method bundle implementing the variant to call.
+  public let bundle: BundleReference<MethodDecl>
+
+  /// The type of the bundle.
+  public let bundleType: MethodType
+
+  /// The variants of the bundle.
+  public let variants: [AccessEffect: Function.ID]
+
+  /// The return storage and arguments of the call.
+  public private(set) var operands: [Operand]
+
+  /// The site of the code corresponding to that instruction.
+  public let site: SourceRange
+
+  /// Creates an instance with the given properties.
+  fileprivate init(
+    bundle: BundleReference<MethodDecl>,
+    bundleType: MethodType,
+    variants: [AccessEffect: Function.ID],
+    output: Operand,
+    arguments: [Operand],
+    site: SourceRange
+  ) {
+    self.bundle = bundle
+    self.bundleType = bundleType
+    self.variants = variants
+    self.operands = [output] + arguments
+    self.site = site
+  }
+
+  /// The capabilities possibly requested on the receiver.
+  public var capabilities: AccessEffectSet {
+    .init(variants.keys)
+  }
+
+  /// The location at which the result of `callee` is stored.
+  public var output: Operand { operands[0] }
+
+  /// The arguments of the call.
+  public var arguments: ArraySlice<Operand> { operands[1...] }
+
+  public mutating func replaceOperand(at i: Int, with new: Operand) {
+    operands[i] = new
+  }
+
+}
+
+extension CallBundle: CustomStringConvertible {
+
+  public var description: String {
+    "call_bundle \(capabilities) \(bundle)(\(list: arguments)) to \(output)"
+  }
+
+}
+
+extension Module {
+
+  /// Creates a `call_bundle` anchored at `site` that calls one of the variants in `bundle` on
+  /// `arguments`.
+  mutating func makeCallBundle(
+    applying bundle: ScopedValue<BundleReference<MethodDecl>>, to arguments: [Operand],
+    writingResultTo output: Operand,
+    at site: SourceRange
+  ) -> CallBundle {
+    var variants: [AccessEffect: Function.ID] = [:]
+    for v in program[bundle.value.bundle].impls {
+      variants[program[v].introducer.value] = demandDeclaration(lowering: v)
+    }
+
+    let bundleType = program.canonicalType(
+      of: bundle.value.bundle, specializedBy: bundle.value.arguments, in: bundle.scope)
+    let t = MethodType(bundleType)!
+
+    precondition((t.inputs.count + 1) == arguments.count)
+    precondition(arguments.allSatisfy({ self[$0] is Access }))
+    precondition(isBorrowSet(output))
+
+    return .init(
+      bundle: bundle.value, bundleType: t, variants: variants,
+      output: output, arguments: arguments,
+      site: site)
+  }
+
+}