Merge pull request #27640 from ravikandhadai/constexpr-closure-support

[Constant Evaluator] Add support for interpreting SIL code with partial applies (i.e., closure creations).

Author: ravikandhadai

include/swift/SIL/SILConstants.h

@@ -33,6 +33,7 @@ struct DerivedAddressValue;
 struct EnumWithPayloadSymbolicValue;
 struct SymbolicValueMemoryObject;
 struct UnknownSymbolicValue;
+struct SymbolicClosure;
 
 extern llvm::cl::opt<unsigned> ConstExprLimit;
 
@@ -261,6 +262,9 @@ class SymbolicValue {
 
     /// This represents an array.
     RK_Array,
+
+    /// This represents a closure.
+    RK_Closure,
   };
 
   union {
@@ -305,7 +309,7 @@ class SymbolicValue {
     /// information about the memory object and access path of the access.
     DerivedAddressValue *derivedAddress;
 
-    // The following fields are for representing an Array.
+    // The following two fields are for representing an Array.
     //
     // In Swift, an array is a non-trivial struct that stores a reference to an
     // internal storage: _ContiguousArrayStorage. Though arrays have value
@@ -329,6 +333,10 @@ class SymbolicValue {
     /// When this symbolic value is of an "Array" kind, this stores a memory
     /// object that contains a SymbolicArrayStorage value.
     SymbolicValueMemoryObject *array;
+
+    /// When this symbolic value is of "Closure" kind, store a pointer to the
+    /// symbolic representation of the closure.
+    SymbolicClosure *closure;
   } value;
 
   RepresentationKind representationKind : 8;
@@ -384,6 +392,9 @@ class SymbolicValue {
     /// This represents an array value.
     Array,
 
+    /// This represents a closure.
+    Closure,
+
     /// These values are generally only seen internally to the system, external
     /// clients shouldn't have to deal with them.
     UninitMemory
@@ -533,6 +544,22 @@ class SymbolicValue {
   /// Return the type of this array symbolic value.
   Type getArrayType() const;
 
+  /// Create and return a symbolic value that represents a closure.
+  /// \param target SILFunction corresponding the target of the closure.
+  /// \param capturedArguments an array consisting of SILValues of captured
+  /// arguments along with their symbolic values when available.
+  /// \param allocator the allocator to use for storing the contents of this
+  /// symbolic value.
+  static SymbolicValue makeClosure(
+      SILFunction *target,
+      ArrayRef<std::pair<SILValue, Optional<SymbolicValue>>> capturedArguments,
+      SymbolicValueAllocator &allocator);
+
+  SymbolicClosure *getClosure() const {
+    assert(getKind() == Closure);
+    return value.closure;
+  }
+
   //===--------------------------------------------------------------------===//
   // Helpers
 
@@ -607,6 +634,57 @@ struct SymbolicValueMemoryObject {
   SymbolicValueMemoryObject(const SymbolicValueMemoryObject &) = delete;
   void operator=(const SymbolicValueMemoryObject &) = delete;
 };
+
+using SymbolicClosureArgument = std::pair<SILValue, Optional<SymbolicValue>>;
+
+/// Representation of a symbolic closure. A symbolic closure consists of a
+/// SILFunction and an array of SIL values, corresponding to the captured
+/// arguments, and (optional) symbolic values representing the constant values
+/// of the captured arguments. The symbolic values are optional as it is not
+/// necessary for every captured argument to be a constant, which enables
+/// representing closures whose captured arguments are not compile-time
+/// constants.
+struct SymbolicClosure final
+  : private llvm::TrailingObjects<SymbolicClosure, SymbolicClosureArgument> {
+
+  friend class llvm::TrailingObjects<SymbolicClosure, SymbolicClosureArgument>;
+
+private:
+
+  SILFunction *target;
+
+  // The number of SIL values captured by the closure.
+  unsigned numCaptures;
+
+  // True iff there exists captured arguments whose constant value is not known.
+  bool hasNonConstantCaptures = true;
+
+  SymbolicClosure() = delete;
+  SymbolicClosure(const SymbolicClosure &) = delete;
+  SymbolicClosure(SILFunction *callee, unsigned numArguments,
+                  bool nonConstantCaptures)
+      : target(callee), numCaptures(numArguments),
+        hasNonConstantCaptures(nonConstantCaptures) {}
+
+public:
+  static SymbolicClosure *create(SILFunction *callee,
+                                 ArrayRef<SymbolicClosureArgument> args,
+                                 SymbolicValueAllocator &allocator);
+
+  ArrayRef<SymbolicClosureArgument> getCaptures() const {
+    return {getTrailingObjects<SymbolicClosureArgument>(), numCaptures};
+  }
+
+  // This is used by the llvm::TrailingObjects base class.
+  size_t numTrailingObjects(OverloadToken<SymbolicClosureArgument>) const {
+    return numCaptures;
+  }
+
+  SILFunction *getTarget() {
+    return target;
+  }
+};
+
 } // end namespace swift
 
 #endif

lib/SIL/SILConstants.cpp

@@ -126,6 +126,29 @@ void SymbolicValue::print(llvm::raw_ostream &os, unsigned indent) const {
   case RK_Array: {
     os << getArrayType() << ": \n";
     getStorageOfArray().print(os, indent);
+    return;
+  }
+  case RK_Closure: {
+    SymbolicClosure *clo = getClosure();
+    SILFunction *target = clo->getTarget();
+    std::string targetName = target->getName();
+    os << "closure: target: " << targetName;
+    ArrayRef<SymbolicClosureArgument> args = clo->getCaptures();
+    os << " captures [\n";
+    for (SymbolicClosureArgument closureArg : args) {
+      os.indent(indent + 2) << closureArg.first << "\n";
+    }
+    os.indent(indent) << "] values: [\n";
+    for (SymbolicClosureArgument closureArg : args) {
+      Optional<SymbolicValue> value = closureArg.second;
+      if (!value.hasValue()) {
+        os.indent(indent + 2) << "nil\n";
+        continue;
+      }
+      value->print(os, indent + 2);
+    }
+    os.indent(indent) << "]\n";
+    return;
   }
   }
 }
@@ -162,6 +185,8 @@ SymbolicValue::Kind SymbolicValue::getKind() const {
     return ArrayStorage;
   case RK_Array:
     return Array;
+  case RK_Closure:
+    return Closure;
   }
   llvm_unreachable("covered switch");
 }
@@ -219,6 +244,11 @@ SymbolicValue::cloneInto(SymbolicValueAllocator &allocator) const {
     SymbolicValue clonedStorage = getStorageOfArray().cloneInto(allocator);
     return getArray(getArrayType(), clonedStorage, allocator);
   }
+  case RK_Closure: {
+    SymbolicClosure *clo = getClosure();
+    ArrayRef<SymbolicClosureArgument> closureArgs = clo->getCaptures();
+    return SymbolicValue::makeClosure(clo->getTarget(), closureArgs, allocator);
+  }
   }
   llvm_unreachable("covered switch");
 }
@@ -661,6 +691,44 @@ Type SymbolicValue::getArrayType() const {
   return value.array->getType();
 }
 
+//===----------------------------------------------------------------------===//
+// Symbolic Closure
+//===----------------------------------------------------------------------===//
+
+SymbolicValue SymbolicValue::makeClosure(SILFunction *target,
+                                         ArrayRef<SymbolicClosureArgument> args,
+                                         SymbolicValueAllocator &allocator) {
+  auto clo = SymbolicClosure::create(target, args, allocator);
+  SymbolicValue result;
+  result.representationKind = RK_Closure;
+  result.value.closure = clo;
+  return result;
+}
+
+SymbolicClosure *SymbolicClosure::create(SILFunction *target,
+                                         ArrayRef<SymbolicClosureArgument> args,
+                                         SymbolicValueAllocator &allocator) {
+  // Determine whether there are captured arguments without a symbolic value.
+  bool hasNonConstantCapture = false;
+  for (SymbolicClosureArgument closureArg : args) {
+    if (!closureArg.second) {
+      hasNonConstantCapture = true;
+      break;
+    }
+  }
+
+  auto byteSizeOfArgs =
+      SymbolicClosure::totalSizeToAlloc<SymbolicClosureArgument>(args.size());
+  auto rawMem = allocator.allocate(byteSizeOfArgs, alignof(SymbolicClosure));
+  //  Placement initialize the object.
+  auto closure = ::new (rawMem)
+      SymbolicClosure(target, args.size(), hasNonConstantCapture);
+  std::uninitialized_copy(
+      args.begin(), args.end(),
+      closure->getTrailingObjects<SymbolicClosureArgument>());
+  return closure;
+}
+
 //===----------------------------------------------------------------------===//
 // Higher level code
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Utils/ConstExpr.cpp

@@ -1630,6 +1630,36 @@ ConstExprFunctionState::evaluateFlowSensitive(SILInstruction *inst) {
                           injectEnumInst->getOperand());
   }
 
+  if (auto *papply = dyn_cast<PartialApplyInst>(inst)) {
+    SILValue calleeOperand = papply->getOperand(0);
+    SymbolicValue calleeValue = getConstantValue(calleeOperand);
+    if (!calleeValue.isConstant())
+      return calleeValue;
+    if (calleeValue.getKind() != SymbolicValue::Function) {
+      return getUnknown(evaluator, (SILInstruction *)papply,
+                        UnknownReason::InvalidOperandValue);
+    }
+
+    SILFunction *target = calleeValue.getFunctionValue();
+    assert(target != nullptr);
+
+    // Arguments to this partial-apply instruction are the captures of the
+    // closure.
+    SmallVector<SymbolicClosureArgument, 4> captures;
+    for (SILValue argument : papply->getArguments()) {
+      SymbolicValue argumentValue = getConstantValue(argument);
+      if (!argumentValue.isConstant()) {
+        captures.push_back({ argument, None });
+        continue;
+      }
+      captures.push_back({ argument, argumentValue });
+    }
+    auto closureVal = SymbolicValue::makeClosure(target, captures,
+                                                 evaluator.getAllocator());
+    setValue(papply, closureVal);
+    return None;
+  }
+
   // If the instruction produces a result, try computing it, and fail if the
   // computation fails.
   if (auto *singleValueInst = dyn_cast<SingleValueInstruction>(inst)) {
@@ -1934,6 +1964,8 @@ ConstExprStepEvaluator::skipByMakingEffectsNonConstant(
            constKind == SymbolicValue::Aggregate ||
            constKind == SymbolicValue::Enum ||
            constKind == SymbolicValue::EnumWithPayload ||
+           constKind == SymbolicValue::Array ||
+           constKind == SymbolicValue::Closure ||
            constKind == SymbolicValue::UninitMemory);
 
     if (constKind != SymbolicValue::Address) {

test/SILOptimizer/constant_evaluable_subset_test.swift

@@ -784,3 +784,52 @@ func testArrayAppendNonEmpty(_ x: String) -> [String] {
 func interpretArrayAppendNonEmpty() -> [String] {
   return testArrayAppendNonEmpty("mkdir")
 }
+
+// CHECK-LABEL: @testClosureInit
+// CHECK-NOT: error:
+@_semantics("constant_evaluable")
+func testClosureInit(_ x: Int) -> () -> Int {
+  return { x }
+}
+
+@_semantics("test_driver")
+func interpretClosureCreation() -> () -> Int {
+  return testClosureInit(19)
+}
+
+// CHECK-LABEL: @testClosureChaining
+// CHECK-NOT: error:
+@_semantics("constant_evaluable")
+func testClosureChaining(_ x: Int, _ y: Int) -> () -> Int {
+  let clo: (Int) -> Int = { $0 + x }
+  return { clo(y) }
+}
+
+@_semantics("test_driver")
+func interpretClosureChains() -> () -> Int {
+  return testClosureChaining(191, 201)
+}
+
+// CHECK-LABEL: @testClosureWithNonConstantCaptures
+// CHECK-NOT: error:
+@_semantics("constant_evaluable")
+func testClosureWithNonConstantCaptures(_ x: @escaping () -> Int) -> () -> Int {
+  return x
+}
+
+@_semantics("test_driver")
+func interpretClosureWithNonConstantCaptures(_ x: Int) -> () -> Int {
+  return testClosureWithNonConstantCaptures({ x })
+}
+
+// CHECK-LABEL: @testAutoClosure
+// CHECK-NOT: error:
+@_semantics("constant_evaluable")
+func testAutoClosure(_ x: @escaping @autoclosure () -> Int) -> () -> Int {
+  return x
+}
+
+@_semantics("test_driver")
+func interpretAutoClosure(_ x: Int) -> () -> Int {
+  return testAutoClosure(x)
+}