[Constant Evaluator] Add support for interpreting SIL code with

partial applies (i.e., closure creations).

A new SymbolicValue: SymbolicClosure represents a closure. It
tracks the SILFunction corresponding to the target of the closure
and the SIL and Symbolic Values of the captured arguments. The
representation does not impose that all captured values must have an
associated symbolic value. This allows the evaluator to track
creations of closures whose captured arguments or bodies are not
constant evaluable.

This commit does not add support for closure applications. It only
adds suppport for closure creations, which correspond to partial_apply
instructions in SIL.

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)
+}