Merge pull request #21869 from apple/marcrasi-const-evaluator-enums

Author: swift-ci

include/swift/SIL/SILConstants.h

@@ -103,6 +103,12 @@ class SymbolicValue {
     /// "aggregate" member of the value union.
     RK_Aggregate,
 
+    /// This value is an enum with no payload.
+    RK_Enum,
+
+    /// This value is an enum with a payload.
+    RK_EnumWithPayload,
+
     /// This represents the address of a memory object.
     RK_DirectAddress,
 
@@ -136,6 +142,14 @@ class SymbolicValue {
     /// information about the array elements and count.
     const SymbolicValue *aggregate;
 
+    /// When this SymbolicValue is of "Enum" kind, this pointer stores
+    /// information about the enum case type.
+    EnumElementDecl *enumVal;
+
+    /// When this SymbolicValue is of "EnumWithPayload" kind, this pointer
+    /// stores information about the enum case type and its payload.
+    EnumWithPayloadSymbolicValue *enumValWithPayload;
+
     /// When the representationKind is "DirectAddress", this pointer is the
     /// memory object referenced.
     SymbolicValueMemoryObject *directAddress;
@@ -186,6 +200,12 @@ class SymbolicValue {
     /// This can be an array, struct, tuple, etc.
     Aggregate,
 
+    /// This is an enum without payload.
+    Enum,
+
+    /// This is an enum with payload (formally known as "associated value").
+    EnumWithPayload,
+
     /// This value represents the address of, or into, a memory object.
     Address,
 
@@ -271,6 +291,25 @@ class SymbolicValue {
 
   ArrayRef<SymbolicValue> getAggregateValue() const;
 
+  /// This returns a constant Symbolic value for the enum case in `decl`, which
+  /// must not have an associated value.
+  static SymbolicValue getEnum(EnumElementDecl *decl) {
+    assert(decl);
+    SymbolicValue result;
+    result.representationKind = RK_Enum;
+    result.value.enumVal = decl;
+    return result;
+  }
+
+  /// `payload` must be a constant.
+  static SymbolicValue getEnumWithPayload(EnumElementDecl *decl,
+                                          SymbolicValue payload,
+                                          ASTContext &astContext);
+
+  EnumElementDecl *getEnumValue() const;
+
+  SymbolicValue getEnumPayloadValue() const;
+
   /// Return a symbolic value that represents the address of a memory object.
   static SymbolicValue getAddress(SymbolicValueMemoryObject *memoryObject) {
     SymbolicValue result;

lib/SIL/SILConstants.cpp

@@ -82,6 +82,20 @@ void SymbolicValue::print(llvm::raw_ostream &os, unsigned indent) const {
       return;
     }
   }
+  case RK_Enum: {
+    auto *decl = getEnumValue();
+    os << "enum: ";
+    decl->print(os);
+    return;
+  }
+  case RK_EnumWithPayload: {
+    auto *decl = getEnumValue();
+    os << "enum: ";
+    decl->print(os);
+    os << ", payload: ";
+    getEnumPayloadValue().print(os, indent);
+    return;
+  }
   case RK_DirectAddress:
   case RK_DerivedAddress: {
     SmallVector<unsigned, 4> accessPath;
@@ -111,6 +125,10 @@ SymbolicValue::Kind SymbolicValue::getKind() const {
     return Function;
   case RK_Aggregate:
     return Aggregate;
+  case RK_Enum:
+    return Enum;
+  case RK_EnumWithPayload:
+    return EnumWithPayload;
   case RK_Integer:
   case RK_IntegerInline:
     return Integer;
@@ -133,6 +151,9 @@ SymbolicValue::cloneInto(ASTContext &astContext) const {
   case RK_Metatype:
   case RK_Function:
     assert(0 && "cloning this representation kind is not supported");
+  case RK_Enum:
+    // These have trivial inline storage, just return a copy.
+    return *this;
   case RK_IntegerInline:
   case RK_Integer:
     return SymbolicValue::getInteger(getIntegerValue(), astContext);
@@ -146,6 +167,8 @@ SymbolicValue::cloneInto(ASTContext &astContext) const {
       results.push_back(elt.cloneInto(astContext));
     return getAggregate(results, astContext);
   }
+  case RK_EnumWithPayload:
+    return getEnumWithPayload(getEnumValue(), getEnumPayloadValue(), astContext);
   case RK_DirectAddress:
   case RK_DerivedAddress: {
     SmallVector<unsigned, 4> accessPath;
@@ -354,6 +377,56 @@ UnknownReason SymbolicValue::getUnknownReason() const {
   return value.unknown->reason;
 }
 
+//===----------------------------------------------------------------------===//
+// Enums
+//===----------------------------------------------------------------------===//
+
+namespace swift {
+
+/// This is the representation of a constant enum value with payload.
+struct EnumWithPayloadSymbolicValue final {
+  /// The enum case.
+  EnumElementDecl *enumDecl;
+  SymbolicValue payload;
+
+  EnumWithPayloadSymbolicValue(EnumElementDecl *decl, SymbolicValue payload)
+      : enumDecl(decl), payload(payload) {}
+
+private:
+  EnumWithPayloadSymbolicValue() = delete;
+  EnumWithPayloadSymbolicValue(const EnumWithPayloadSymbolicValue &) = delete;
+};
+} // end namespace swift
+
+/// This returns a constant Symbolic value for the enum case in `decl` with a
+/// payload.
+SymbolicValue
+SymbolicValue::getEnumWithPayload(EnumElementDecl *decl, SymbolicValue payload,
+                                  ASTContext &astContext) {
+  assert(decl && payload.isConstant());
+  auto rawMem = astContext.Allocate(sizeof(EnumWithPayloadSymbolicValue),
+                                    alignof(EnumWithPayloadSymbolicValue));
+  auto enumVal = ::new (rawMem) EnumWithPayloadSymbolicValue(decl, payload);
+
+  SymbolicValue result;
+  result.representationKind = RK_EnumWithPayload;
+  result.value.enumValWithPayload = enumVal;
+  return result;
+}
+
+EnumElementDecl *SymbolicValue::getEnumValue() const {
+  if (representationKind == RK_Enum)
+    return value.enumVal;
+
+  assert(representationKind == RK_EnumWithPayload);
+  return value.enumValWithPayload->enumDecl;
+}
+
+SymbolicValue SymbolicValue::getEnumPayloadValue() const {
+  assert(representationKind == RK_EnumWithPayload);
+  return value.enumValWithPayload->payload;
+}
+
 //===----------------------------------------------------------------------===//
 // Addresses
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Utils/ConstExpr.cpp

@@ -287,6 +287,25 @@ SymbolicValue ConstExprFunctionState::computeConstantValue(SILValue value) {
     return calculatedValues[apply];
   }
 
+  if (auto *enumVal = dyn_cast<EnumInst>(value)) {
+    if (!enumVal->hasOperand())
+      return SymbolicValue::getEnum(enumVal->getElement());
+
+    auto payload = getConstantValue(enumVal->getOperand());
+    if (!payload.isConstant())
+      return payload;
+    return SymbolicValue::getEnumWithPayload(enumVal->getElement(), payload,
+                                             evaluator.getASTContext());
+  }
+
+  // This one returns the address of its enum payload.
+  if (auto *dai = dyn_cast<UncheckedTakeEnumDataAddrInst>(value)) {
+    auto enumVal = getConstAddrAndLoadResult(dai->getOperand());
+    if (!enumVal.isConstant())
+      return enumVal;
+    return createMemoryObject(value, enumVal.getEnumPayloadValue());
+  }
+
   // This instruction is a marker that returns its first operand.
   if (auto *bai = dyn_cast<BeginAccessInst>(value))
     return getConstantValue(bai->getOperand());
@@ -1244,6 +1263,34 @@ static llvm::Optional<SymbolicValue> evaluateAndCacheCall(
       continue;
     }
 
+    if (isa<SwitchEnumAddrInst>(inst) || isa<SwitchEnumInst>(inst)) {
+      SymbolicValue value;
+      SwitchEnumInstBase *switchInst = dyn_cast<SwitchEnumInst>(inst);
+      if (switchInst) {
+        value = state.getConstantValue(switchInst->getOperand());
+      } else {
+        switchInst = cast<SwitchEnumAddrInst>(inst);
+        value = state.getConstAddrAndLoadResult(switchInst->getOperand());
+      }
+      if (!value.isConstant())
+        return value;
+      assert(value.getKind() == SymbolicValue::Enum ||
+             value.getKind() == SymbolicValue::EnumWithPayload);
+      // Set up basic block arguments.
+      auto *caseBB = switchInst->getCaseDestination(value.getEnumValue());
+      if (caseBB->getNumArguments() > 0) {
+        assert(value.getKind() == SymbolicValue::EnumWithPayload);
+        // When there are multiple payload components, they form a single
+        // tuple-typed argument.
+        assert(caseBB->getNumArguments() == 1);
+        auto argument = value.getEnumPayloadValue();
+        assert(argument.isConstant());
+        state.setValue(caseBB->getArgument(0), argument);
+      }
+      nextInst = caseBB->begin();
+      continue;
+    }
+
     LLVM_DEBUG(llvm::dbgs()
                << "ConstExpr: Unknown Terminator: " << *inst << "\n");
 

test/SILOptimizer/pound_assert.swift

@@ -497,3 +497,91 @@ func stringInitNonEmptyFlowSensitive() -> ContainsString {
 func invokeStringInitNonEmptyFlowSensitive() {
   #assert(stringInitNonEmptyFlowSensitive().x == 1)
 }
+
+//===----------------------------------------------------------------------===//
+// Enums and optionals.
+//===----------------------------------------------------------------------===//
+func isNil(_ x: Int?) -> Bool {
+  return x == nil
+}
+
+#assert(isNil(nil))
+#assert(!isNil(3))
+
+public enum Pet {
+  case bird
+  case cat(Int)
+  case dog(Int, Int)
+  case fish
+}
+
+public func weighPet(pet: Pet) -> Int {
+  switch pet {
+  case .bird: return 3
+  case let .cat(weight): return weight
+  case let .dog(w1, w2): return w1+w2
+  default: return 1
+  }
+}
+
+#assert(weighPet(pet: .bird) == 3)
+#assert(weighPet(pet: .fish) == 1)
+#assert(weighPet(pet: .cat(2)) == 2)
+// expected-error @+1 {{assertion failed}}
+#assert(weighPet(pet: .cat(2)) == 3)
+#assert(weighPet(pet: .dog(9, 10)) == 19)
+
+// Test indirect enums.
+indirect enum IntExpr {
+  case int(_ value: Int)
+  case add(_ lhs: IntExpr, _ rhs: IntExpr)
+  case multiply(_ lhs: IntExpr, _ rhs: IntExpr)
+}
+
+func evaluate(intExpr: IntExpr) -> Int {
+  switch intExpr {
+  case .int(let value):
+    return value
+  case .add(let lhs, let rhs):
+    return evaluate(intExpr: lhs) + evaluate(intExpr: rhs)
+  case .multiply(let lhs, let rhs):
+    return evaluate(intExpr: lhs) * evaluate(intExpr: rhs)
+  }
+}
+
+// TODO: The constant evaluator can't handle indirect enums yet.
+// expected-error @+2 {{#assert condition not constant}}
+// expected-note @+1 {{could not fold operation}}
+#assert(evaluate(intExpr: .int(5)) == 5)
+// expected-error @+2 {{#assert condition not constant}}
+// expected-note @+1 {{could not fold operation}}
+#assert(evaluate(intExpr: .add(.int(5), .int(6))) == 11)
+// expected-error @+2 {{#assert condition not constant}}
+// expected-note @+1 {{could not fold operation}}
+#assert(evaluate(intExpr: .add(.multiply(.int(2), .int(2)), .int(3))) == 7)
+
+// Test address-only enums.
+protocol IntContainerProtocol {
+  var value: Int { get }
+}
+
+struct IntContainer : IntContainerProtocol {
+  let value: Int
+}
+
+enum AddressOnlyEnum<T: IntContainerProtocol> {
+  case double(_ value: T)
+  case triple(_ value: T)
+}
+
+func evaluate<T>(addressOnlyEnum: AddressOnlyEnum<T>) -> Int {
+  switch addressOnlyEnum {
+  case .double(let value):
+    return 2 * value.value
+  case .triple(let value):
+    return 3 * value.value
+  }
+}
+
+#assert(evaluate(addressOnlyEnum: .double(IntContainer(value: 1))) == 2)
+#assert(evaluate(addressOnlyEnum: .triple(IntContainer(value: 1))) == 3)