[embedded] Implement non-allocating embedded Swift mode, under -no-allocations flag

include/swift/AST/DiagnosticsFrontend.def

@@ -561,6 +561,8 @@ ERROR(wmo_with_embedded,none,
       "Whole module optimization (wmo) must be enabled with embedded Swift.", ())
 ERROR(objc_with_embedded,none,
       "Objective-C interoperability cannot be enabled with embedded Swift.", ())
+ERROR(no_allocations_without_embedded,none,
+      "-no-allocations is only applicable with embedded Swift.", ())
 
 #define UNDEFINE_DIAGNOSTIC_MACROS
 #include "DefineDiagnosticMacros.h"

include/swift/AST/DiagnosticsSIL.def

@@ -364,6 +364,10 @@ ERROR(embedded_swift_metatype_type,none,
       "cannot use metatype of type %0 in embedded Swift", (Type))
 ERROR(embedded_swift_metatype,none,
       "cannot use metatype in embedded Swift", ())
+ERROR(embedded_swift_allocating_type,none,
+      "cannot use allocating type %0 in -no-allocations mode", (Type))
+ERROR(embedded_swift_allocating,none,
+      "cannot use allocating type in -no-allocations mode", ())
 NOTE(performance_called_from,none,
       "called from here", ())
 

include/swift/AST/SILOptions.h

@@ -282,6 +282,9 @@ class SILOptions {
   /// Are we building in embedded Swift mode?
   bool EmbeddedSwift = false;
 
+  /// Are we building in embedded Swift + -no-allocations?
+  bool NoAllocations = false;
+
   /// The name of the file to which the backend should save optimization
   /// records.
   std::string OptRecordFile;

include/swift/Option/Options.td

@@ -872,6 +872,10 @@ def save_optimization_record_passes :
            "the generated optimization record "
            "(by default, include all passes)">, MetaVarName<"<regex>">;
 
+def no_allocations : Flag<["-"], "no-allocations">,
+  Flags<[FrontendOption, HelpHidden]>,
+  HelpText<"Diagnose any code that needs to heap allocate (classes, closures, etc.)">;
+
 // Platform options.
 def enable_app_extension : Flag<["-"], "application-extension">,
   Flags<[FrontendOption, NoInteractiveOption]>,

lib/Frontend/CompilerInvocation.cpp

@@ -2425,6 +2425,8 @@ static bool ParseSILArgs(SILOptions &Opts, ArgList &Args,
 
   Opts.OSSACompleteLifetimes |= Args.hasArg(OPT_enable_ossa_complete_lifetimes);
 
+  Opts.NoAllocations = Args.hasArg(OPT_no_allocations);
+
   return false;
 }
 
@@ -3228,6 +3230,11 @@ bool CompilerInvocation::parseArgs(
     SILOpts.SkipFunctionBodies = FunctionBodySkipping::None;
     SILOpts.CMOMode = CrossModuleOptimizationMode::Everything;
     SILOpts.EmbeddedSwift = true;
+  } else {
+    if (SILOpts.NoAllocations) {
+      Diags.diagnose(SourceLoc(), diag::no_allocations_without_embedded);
+      return true;
+    }
   }
 
   return false;

lib/SILOptimizer/Mandatory/PerformanceDiagnostics.cpp

@@ -94,6 +94,10 @@ class PerformanceDiagnostics {
     return visitFunction(function, perfConstr, /*parentLoc*/ nullptr);
   }
 
+  bool visitFunctionEmbeddedSwift(SILFunction *function) {
+    return visitFunctionEmbeddedSwift(function, /*parentLoc*/ nullptr);
+  }
+
   /// Check functions _without_ performance annotations.
   ///
   /// This is need to check closure arguments of called performance-annotated
@@ -104,6 +108,9 @@ class PerformanceDiagnostics {
   bool visitFunction(SILFunction *function, PerformanceConstraints perfConstr,
                         LocWithParent *parentLoc);
 
+  bool visitFunctionEmbeddedSwift(SILFunction *function,
+                                  LocWithParent *parentLoc);
+
   bool visitInst(SILInstruction *inst, PerformanceConstraints perfConstr,
                     LocWithParent *parentLoc);
 
@@ -150,6 +157,75 @@ static bool isEffectFreeArraySemanticCall(SILInstruction *inst) {
   }
 }
 
+/// Prints Embedded Swift specific performance diagnostics (no existentials,
+/// no metatypes, optionally no allocations) for \p function.
+bool PerformanceDiagnostics::visitFunctionEmbeddedSwift(
+    SILFunction *function, LocWithParent *parentLoc) {
+  // Don't check generic functions in embedded Swift, they're about to be
+  // removed anyway.
+  if (function->getLoweredFunctionType()->getSubstGenericSignature())
+    return false;
+
+  if (!function->isDefinition())
+    return false;
+
+  if (visitedFuncs.contains(function))
+    return false;
+  visitedFuncs[function] = PerformanceConstraints::None;
+
+  NonErrorHandlingBlocks neBlocks(function);
+
+  for (SILBasicBlock &block : *function) {
+    for (SILInstruction &inst : block) {
+      if (visitInst(&inst, PerformanceConstraints::None, parentLoc)) {
+        if (inst.getLoc().getSourceLoc().isInvalid()) {
+          auto demangledName = Demangle::demangleSymbolAsString(
+              inst.getFunction()->getName(),
+              Demangle::DemangleOptions::SimplifiedUIDemangleOptions());
+          llvm::errs() << "in function " << demangledName << "\n";
+        }
+        LLVM_DEBUG(llvm::dbgs() << inst << *inst.getFunction());
+        return true;
+      }
+
+      if (auto as = FullApplySite::isa(&inst)) {
+        LocWithParent asLoc(inst.getLoc().getSourceLoc(), parentLoc);
+        LocWithParent *loc = &asLoc;
+        if (parentLoc &&
+            asLoc.loc == inst.getFunction()->getLocation().getSourceLoc())
+          loc = parentLoc;
+
+        for (SILFunction *callee : bca->getCalleeList(as)) {
+          if (visitFunctionEmbeddedSwift(callee, loc))
+            return true;
+        }
+      } else if (auto *bi = dyn_cast<BuiltinInst>(&inst)) {
+        PrettyStackTracePerformanceDiagnostics stackTrace(
+            "visitFunction::BuiltinInst (once, once with context)", &inst);
+
+        switch (bi->getBuiltinInfo().ID) {
+        case BuiltinValueKind::Once:
+        case BuiltinValueKind::OnceWithContext:
+          if (auto *fri = dyn_cast<FunctionRefInst>(bi->getArguments()[1])) {
+            LocWithParent asLoc(bi->getLoc().getSourceLoc(), parentLoc);
+            LocWithParent *loc = &asLoc;
+            if (parentLoc &&
+                asLoc.loc == bi->getFunction()->getLocation().getSourceLoc())
+              loc = parentLoc;
+
+            if (visitFunctionEmbeddedSwift(fri->getReferencedFunction(), loc))
+              return true;
+          }
+          break;
+        default:
+          break;
+        }
+      }
+    }
+  }
+  return false;
+}
+
 /// Prints performance diagnostics for \p function.
 bool PerformanceDiagnostics::visitFunction(SILFunction *function,
                                               PerformanceConstraints perfConstr,
@@ -438,6 +514,18 @@ bool PerformanceDiagnostics::visitInst(SILInstruction *inst,
         return true;
       }
     }
+
+    if (module.getOptions().NoAllocations) {
+      if (impact & RuntimeEffect::Allocating) {
+        PrettyStackTracePerformanceDiagnostics stackTrace("allocation", inst);
+        if (impactType) {
+          diagnose(loc, diag::embedded_swift_allocating_type, impactType.getASTType());
+        } else {
+          diagnose(loc, diag::embedded_swift_allocating);
+        }
+        return true;
+      }
+    }
   }
 
   if (perfConstr == PerformanceConstraints::None ||
@@ -585,19 +673,6 @@ bool PerformanceDiagnostics::visitInst(SILInstruction *inst,
 void PerformanceDiagnostics::checkNonAnnotatedFunction(SILFunction *function) {
   for (SILBasicBlock &block : *function) {
     for (SILInstruction &inst : block) {
-      if (function->getModule().getOptions().EmbeddedSwift) {
-        auto loc = LocWithParent(inst.getLoc().getSourceLoc(), nullptr);
-        if (visitInst(&inst, PerformanceConstraints::None, &loc)) {
-          if (inst.getLoc().getSourceLoc().isInvalid()) {
-            auto demangledName = Demangle::demangleSymbolAsString(
-                inst.getFunction()->getName(),
-                Demangle::DemangleOptions::SimplifiedUIDemangleOptions());
-            llvm::errs() << "in function " << demangledName << "\n";
-          }
-          LLVM_DEBUG(llvm::dbgs() << inst << *inst.getFunction());
-        }
-      }
-
       auto as = FullApplySite::isa(&inst);
       if (!as)
         continue;
@@ -701,16 +776,54 @@ class PerformanceDiagnosticsPass : public SILModuleTransform {
       if (function.wasDeserializedCanonical())
         continue;
 
-      // Don't check generic functions in embedded Swift, they're about to be
-      // removed anyway.
-      if (getModule()->getOptions().EmbeddedSwift &&
-          function.getLoweredFunctionType()->getSubstGenericSignature())
-        continue;
-
       if (function.getPerfConstraints() == PerformanceConstraints::None) {
         diagnoser.checkNonAnnotatedFunction(&function);
       }
     }
+
+    if (getModule()->getOptions().EmbeddedSwift) {
+      // Run embedded Swift SIL checks for metatype/existential use, and
+      // allocation use (under -no-allocations mode). Try to start with public
+      // and exported functions to get better call tree information.
+      SmallVector<SILFunction *, 8> externallyVisibleFunctions;
+      SmallVector<SILFunction *, 8> vtableMembers;
+      SmallVector<SILFunction *, 8> others;
+      SmallVector<SILFunction *, 8> constructorsAndDestructors;
+
+      for (SILFunction &function : *module) {
+        auto func = function.getLocation().getAsASTNode<AbstractFunctionDecl>();
+        if (func) {
+          if (isa<DestructorDecl>(func) || isa<ConstructorDecl>(func)) {
+            constructorsAndDestructors.push_back(&function);
+            continue;
+          }
+          if (getMethodDispatch(func) == MethodDispatch::Class) {
+            vtableMembers.push_back(&function);
+            continue;
+          }
+        }
+
+        if (function.isPossiblyUsedExternally()) {
+          externallyVisibleFunctions.push_back(&function);
+          continue;
+        }
+
+        others.push_back(&function);
+      }
+
+      for (SILFunction *function : externallyVisibleFunctions) {
+        diagnoser.visitFunctionEmbeddedSwift(function);
+      }
+      for (SILFunction *function : vtableMembers) {
+        diagnoser.visitFunctionEmbeddedSwift(function);
+      }
+      for (SILFunction *function : others) {
+        diagnoser.visitFunctionEmbeddedSwift(function);
+      }
+      for (SILFunction *function : constructorsAndDestructors) {
+        diagnoser.visitFunctionEmbeddedSwift(function);
+      }
+    }    
   }
 };
 

test/embedded/anyobject-error-no-stdlib.swift

@@ -16,5 +16,4 @@ precedencegroup AssignmentPrecedence { assignment: true }
 
 public func foo(_ x: AnyObject) {
   _ = type(of: x) // expected-error {{cannot use a value of protocol type 'AnyObject' in embedded Swift}}
-  // expected-note@-1 {{called from here}}
 }

test/embedded/basic-errors-no-stdlib.swift

@@ -7,5 +7,4 @@ struct Concrete: Player {}
 
 public func test() -> any Player {
   Concrete() // expected-error {{cannot use a value of protocol type 'any Player' in embedded Swift}}
-  // expected-note@-1 {{called from here}}
 }

test/embedded/metatypes.swift

@@ -9,6 +9,5 @@ public func sink<T>(t: T) {}
 public func test() -> Int {
   let metatype = Int.self
   sink(t: metatype) // expected-error {{cannot use metatype of type 'Int' in embedded Swift}}
-  // expected-note@-1 {{called from here}}
   return 42
 }

test/embedded/no-allocations.swift

@@ -0,0 +1,17 @@
+// RUN: %target-swift-emit-ir %s -wmo
+// RUN: %target-swift-emit-ir %s -enable-experimental-feature Embedded -wmo
+// RUN: %target-swift-emit-ir %s -enable-experimental-feature Embedded -no-allocations -wmo -verify -verify-ignore-unknown
+
+// REQUIRES: swift_in_compiler
+// REQUIRES: OS=macosx || OS=linux-gnu
+
+public class X {} // expected-error {{cannot use allocating type 'X' in -no-allocations mode}}
+public func use_a_class() -> X {
+	let x = X() // expected-note {{called from here}}
+	return x
+}
+
+public func use_an_array() -> Int {
+	let a = [1, 2, 3] // expected-error {{cannot use allocating type '_ContiguousArrayStorage<Int>' in -no-allocations mode}}
+	return a.count
+}

test/embedded/typeof.swift

@@ -7,7 +7,6 @@
 public func unsafeWriteArray<T, R>(_ elementType: R.Type, array: inout T, index n: Int, value: R) {
   precondition(_isPOD(elementType))
   precondition(_isPOD(type(of: array))) // expected-error {{cannot use metatype of type '(Int, Int, Int, Int)' in embedded Swift}}
-  // expected-note@-1 {{called from here}}
 
   return withUnsafeMutableBytes(of: &array) { ptr in
     let buffer = ptr.bindMemory(to: R.self)