StringOptimization: handle static let variables for String constant folding.

For example, constant fold:

struct Str {
    static let s = "hello"
}
...
let x = "<\(Str.s)>"

Author: eeckstein

lib/SILOptimizer/Transforms/StringOptimization.cpp

@@ -19,6 +19,7 @@
 #include "swift/SILOptimizer/Analysis/ValueTracking.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILBasicBlock.h"
+#include "swift/SIL/SILGlobalVariable.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/AST/SemanticAttrs.h"
 #include "swift/AST/ParameterList.h"
@@ -58,13 +59,15 @@ class StringOptimization {
     StringRef str;
 
     /// Negative means: not constant
-    int numCodeUnits = -1;
-    
-    /// Not 0 for the empty-string initializer which reserves a capacity.
     int reservedCapacity = 0;
 
-    bool isConstant() const { return numCodeUnits >= 0; }
+    StringInfo(StringRef str, int reservedCapacity = 0) :
+      str(str), reservedCapacity(reservedCapacity) { }
+    
+    bool isConstant() const { return reservedCapacity >= 0; }
     bool isEmpty() const { return isConstant() && str.empty(); }
+    
+    static StringInfo unknown() { return StringInfo(StringRef(), -1); }
   };
 
   /// The stdlib's String type.
@@ -96,6 +99,8 @@ class StringOptimization {
   static void invalidateModifiedObjects(SILInstruction *inst,
                             llvm::DenseMap<SILValue, SILValue> &storedStrings);
   static StringInfo getStringInfo(SILValue value);
+  static StringInfo getStringFromStaticLet(SILValue value);
+
   static Optional<int> getIntConstant(SILValue value);
   static void replaceAppendWith(ApplyInst *appendCall, SILValue newValue,
                                 bool copyNewValue);
@@ -368,32 +373,31 @@ void StringOptimization::invalidateModifiedObjects(SILInstruction *inst,
 
 /// Returns information about value if it's a constant string.
 StringOptimization::StringInfo StringOptimization::getStringInfo(SILValue value) {
-  // Start with a non-constant result.
-  StringInfo result;
+  if (!value)
+    return StringInfo::unknown();
 
-  auto *apply = dyn_cast_or_null<ApplyInst>(value);
-  if (!apply)
-    return result;
+  auto *apply = dyn_cast<ApplyInst>(value);
+  if (!apply) {
+    return getStringFromStaticLet(value);
+  }
 
   SILFunction *callee = apply->getReferencedFunctionOrNull();
   if (!callee)
-    return result;
+    return StringInfo::unknown();
 
   if (callee->hasSemanticsAttr(semantics::STRING_INIT_EMPTY)) {
     // An empty string initializer.
-    result.numCodeUnits = 0;
-    return result;
+    return StringInfo("");
   }
 
   if (callee->hasSemanticsAttr(semantics::STRING_INIT_EMPTY_WITH_CAPACITY)) {
     // An empty string initializer with initial capacity.
-    result.numCodeUnits = 0;
-    result.reservedCapacity = std::numeric_limits<int>::max();
+    int reservedCapacity = std::numeric_limits<int>::max();
     if (apply->getNumArguments() > 0) {
       if (Optional<int> capacity = getIntConstant(apply->getArgument(0)))
-        result.reservedCapacity = capacity.getValue();
+        reservedCapacity = capacity.getValue();
     }
-    return result;
+    return StringInfo("", reservedCapacity);
   }
 
   if (callee->hasSemanticsAttr(semantics::STRING_MAKE_UTF8)) {
@@ -404,13 +408,79 @@ StringOptimization::StringInfo StringOptimization::getStringInfo(SILValue value)
     auto *intLiteral = dyn_cast<IntegerLiteralInst>(lengthVal);
     if (intLiteral && stringLiteral &&
         // For simplicity, we only support UTF8 string literals.
-        stringLiteral->getEncoding() == StringLiteralInst::Encoding::UTF8) {
-      result.str = stringLiteral->getValue();
-      result.numCodeUnits = intLiteral->getValue().getSExtValue();
-      return result;
+        stringLiteral->getEncoding() == StringLiteralInst::Encoding::UTF8 &&
+        // This passed number of code units should always match the size of the
+        // string in the string literal. Just to be on the safe side, check it.
+        intLiteral->getValue() == stringLiteral->getValue().size()) {
+      return StringInfo(stringLiteral->getValue());
     }
   }
-  return result;
+  return StringInfo::unknown();
+}
+
+/// Return the string if \p value is a load from a global static let, which is
+/// initialized with a String constant.
+StringOptimization::StringInfo
+StringOptimization::getStringFromStaticLet(SILValue value) {
+  // Match the pattern
+  //   %ptr_to_global = apply %addressor()
+  //   %global_addr = pointer_to_address %ptr_to_global
+  //   %value = load %global_addr
+  auto *load = dyn_cast<LoadInst>(value);
+  if (!load)
+    return StringInfo::unknown();
+ 
+  auto *pta = dyn_cast<PointerToAddressInst>(load->getOperand());
+  if (!pta)
+    return StringInfo::unknown();
+    
+  auto *addressorCall = dyn_cast<ApplyInst>(pta->getOperand());
+  if (!addressorCall)
+    return StringInfo::unknown();
+    
+  SILFunction *addressorFunc = addressorCall->getReferencedFunctionOrNull();
+  if (!addressorFunc)
+    return StringInfo::unknown();
+    
+  // The addressor function has a builtin.once call to the initializer.
+  BuiltinInst *onceCall = nullptr;
+  SILFunction *initializer = findInitializer(addressorFunc, onceCall);
+  if (!initializer)
+    return StringInfo::unknown();
+  
+  if (initializer->size() != 1)
+    return StringInfo::unknown();
+
+  // Match the pattern
+  //   %addr = global_addr @staticStringLet
+  //   ...
+  //   %str = apply %stringInitializer(...)
+  //   store %str to %addr
+  GlobalAddrInst *gAddr = nullptr;
+  for (SILInstruction &inst : initializer->front()) {
+    if (auto *ga = dyn_cast<GlobalAddrInst>(&inst)) {
+      if (gAddr)
+        return StringInfo::unknown();
+      gAddr = ga;
+    }
+  }
+  if (!gAddr || !gAddr->getReferencedGlobal()->isLet())
+    return StringInfo::unknown();
+  
+  Operand *gUse = gAddr->getSingleUse();
+  auto *store = dyn_cast<StoreInst>(gUse->getUser());
+  if (!store || store->getDest() != gAddr)
+    return StringInfo::unknown();
+    
+  SILValue initVal = store->getSrc();
+  
+  // This check is probably not needed, but let's be on the safe side:
+  // it prevents an infinite recursion if the initializer of the global is
+  // itself a load of another global, and so on.
+  if (isa<LoadInst>(initVal))
+    return StringInfo::unknown();
+
+  return getStringInfo(initVal);
 }
 
 /// Returns the constant integer value if \a value is an Int or Bool struct with

test/SILOptimizer/string_optimization.swift

@@ -12,6 +12,8 @@ import Foundation
 
 struct Outer {
   struct Inner { }
+
+  static let staticString = "static"
 }
 
 // More types are tested in test/stdlib/TypeName.swift and
@@ -36,6 +38,15 @@ public func testFoldCompleteInterpolation() -> String {
   return "-\([Int].self) \(s) \("cool")+"
 }
 
+// CHECK-LABEL: sil [noinline] @$s4test0A13FoldStaticLetSSyF
+// CHECK-NOT: apply
+// CHECK-NOT: bb1
+// CHECK: } // end sil function '$s4test0A13FoldStaticLetSSyF'
+@inline(never)
+public func testFoldStaticLet() -> String {
+  return "-\(Outer.staticString)+"
+}
+
 // CHECK-LABEL: sil [noinline] @$s4test0A19UnqualifiedTypeNameSSyF 
 // CHECK-NOT: apply
 // CHECK-NOT: bb1
@@ -92,6 +103,9 @@ printEmbeeded(testTypeNameInterpolation())
 // CHECK-OUTPUT: <-Array<Int> is cool+>
 printEmbeeded(testFoldCompleteInterpolation())
 
+// CHECK-OUTPUT: <-static+>
+printEmbeeded(testFoldStaticLet())
+
 // CHECK-OUTPUT: <Inner>
 printEmbeeded(testUnqualifiedTypeName())