Add support for dehydrated runtime data structures

src/coreclr/nativeaot/BuildIntegration/Microsoft.NETCore.Native.targets

@@ -214,6 +214,7 @@ The .NET Foundation licenses this file to you under the MIT license.
       <IlcArg Condition="$(IlcGenerateMetadataLog) == 'true'" Include="--metadatalog:$(NativeIntermediateOutputPath)%(ManagedBinary.Filename).metadata.csv" />
       <IlcArg Condition="$(TargetArchitecture) != ''" Include="--targetarch:$(TargetArchitecture)" />
       <IlcArg Condition="$(IlcMultiModule) == 'true'" Include="--multifile" />
+      <IlcArg Condition="$(IlcMultiModule) != 'true' and '$(IlcDehydrate)' != 'false'" Include="--dehydrate" />
       <IlcArg Condition="$(Optimize) == 'true'" Include="-O" />
       <IlcArg Condition="$(NativeDebugSymbols) == 'true'" Include="-g" />
       <IlcArg Condition="$(IlcDwarfVersion) == '5'" Include="--gdwarf-5" />

src/coreclr/nativeaot/Common/src/Internal/Runtime/CompilerHelpers/StartupCodeHelpers.cs

@@ -85,17 +85,32 @@ private static unsafe TypeManagerHandle[] CreateTypeManagers(IntPtr osModule, In
                     moduleCount++;
             }
 
-            TypeManagerHandle[] modules = new TypeManagerHandle[moduleCount];
+            // We cannot use the new keyword just yet, so stackalloc the array first
+            TypeManagerHandle* pHandles = stackalloc TypeManagerHandle[moduleCount];
             int moduleIndex = 0;
             for (int i = 0; i < count; i++)
             {
                 if (pModuleHeaders[i] != IntPtr.Zero)
                 {
-                    modules[moduleIndex] = RuntimeImports.RhpCreateTypeManager(osModule, pModuleHeaders[i], pClasslibFunctions, nClasslibFunctions);
-                    moduleIndex++;
+                    TypeManagerHandle handle = RuntimeImports.RhpCreateTypeManager(osModule, pModuleHeaders[i], pClasslibFunctions, nClasslibFunctions);
+
+                    // Rehydrate any dehydrated data structures
+                    IntPtr dehydratedDataSection = RuntimeImports.RhGetModuleSection(
+                        handle, ReadyToRunSectionType.DehydratedData, out int dehydratedDataLength);
+                    if (dehydratedDataSection != IntPtr.Zero)
+                    {
+                        RehydrateData(dehydratedDataSection, dehydratedDataLength);
+                    }
+
+                    pHandles[moduleIndex++] = handle;
                 }
             }
 
+            // Any potentially dehydrated MethodTables got rehydrated, we can safely use `new` now.
+            TypeManagerHandle[] modules = new TypeManagerHandle[moduleCount];
+            for (int i = 0; i < moduleCount; i++)
+                modules[i] = pHandles[i];
+
             return modules;
         }
 
@@ -217,6 +232,49 @@ private static unsafe object[] InitializeStatics(IntPtr gcStaticRegionStart, int
 
             return spine;
         }
+
+        private static unsafe void RehydrateData(IntPtr dehydratedData, int length)
+        {
+            // Destination for the hydrated data is in the first 32-bit relative pointer
+            byte* pDest = (byte*)ReadRelPtr32((void*)dehydratedData);
+
+            // The dehydrated data follows
+            byte* pCurrent = (byte*)dehydratedData + sizeof(int);
+            byte* pEnd = (byte*)dehydratedData + length;
+
+            // Fixup table immediately follows the command stream
+            int* pFixups = (int*)pEnd;
+
+            while (pCurrent < pEnd)
+            {
+                pCurrent = DehydratedDataCommand.Decode(pCurrent, out int command, out int payload);
+                switch (command)
+                {
+                    case DehydratedDataCommand.Copy:
+                        // TODO: can we do any kind of memcpy here?
+                        for (; payload > 0; payload--)
+                            *pDest++ = *pCurrent++;
+                        break;
+                    case DehydratedDataCommand.ZeroFill:
+                        pDest += payload;
+                        break;
+                    case DehydratedDataCommand.PtrReloc:
+                        *(void**)pDest = ReadRelPtr32(pFixups + payload);
+                        pDest += sizeof(void*);
+                        break;
+                    case DehydratedDataCommand.RelPtr32Reloc:
+                        WriteRelPtr32(pDest, ReadRelPtr32(pFixups + payload));
+                        pDest += sizeof(int);
+                        break;
+                }
+            }
+
+            static void* ReadRelPtr32(void* address)
+                => (byte*)address + *(int*)address;
+
+            static void WriteRelPtr32(void* dest, void* value)
+                => *(int*)dest = (int)((byte*)value - (byte*)dest);
+        }
     }
 
     [StructLayout(LayoutKind.Sequential)]

src/coreclr/nativeaot/System.Private.CoreLib/src/System.Private.CoreLib.csproj

@@ -40,6 +40,9 @@
 
   <!-- Sources -->
   <ItemGroup>
+    <Compile Include="$(CompilerCommonPath)\Internal\Runtime\DehydratedData.cs">
+      <Link>Internal\Runtime\DehydratedData.cs</Link>
+    </Compile>
     <Compile Include="$(CompilerCommonPath)\Internal\Runtime\RuntimeConstants.cs">
       <Link>Internal\Runtime\RuntimeConstants.cs</Link>
     </Compile>

src/coreclr/nativeaot/Test.CoreLib/src/Test.CoreLib.csproj

@@ -103,6 +103,9 @@
     <Compile Include="$(CompilerCommonPath)\Internal\Runtime\ModuleHeaders.cs">
       <Link>Internal\Runtime\ModuleHeaders.cs</Link>
     </Compile>
+    <Compile Include="$(CompilerCommonPath)\Internal\Runtime\DehydratedData.cs">
+      <Link>Internal\Runtime\DehydratedData.cs</Link>
+    </Compile>
     <Compile Include="$(CompilerCommonPath)\Internal\Runtime\RuntimeConstants.cs">
       <Link>Internal\Runtime\RuntimeConstants.cs</Link>
     </Compile>

src/coreclr/tools/Common/Compiler/DependencyAnalysis/ObjectNodeSection.cs

@@ -7,7 +7,8 @@ public enum SectionType
     {
         ReadOnly,
         Writeable,
-        Executable
+        Executable,
+        Uninitialized,
     }
 
     /// <summary>
@@ -47,6 +48,7 @@ public bool IsStandardSection
         public static readonly ObjectNodeSection TextSection = new ObjectNodeSection("text", SectionType.Executable);
         public static readonly ObjectNodeSection TLSSection = new ObjectNodeSection("TLS", SectionType.Writeable);
         public static readonly ObjectNodeSection BssSection = new ObjectNodeSection("bss", SectionType.Writeable);
+        public static readonly ObjectNodeSection HydrationTargetSection = new ObjectNodeSection("hydrated", SectionType.Uninitialized);
         public static readonly ObjectNodeSection ManagedCodeWindowsContentSection = new ObjectNodeSection(".managedcode$I", SectionType.Executable);
         public static readonly ObjectNodeSection FoldableManagedCodeWindowsContentSection = new ObjectNodeSection(".managedcode$I", SectionType.Executable);
         public static readonly ObjectNodeSection ManagedCodeUnixContentSection = new ObjectNodeSection("__managedcode", SectionType.Executable);

src/coreclr/tools/Common/Compiler/DependencyAnalysis/Relocation.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.Diagnostics;
 
 namespace ILCompiler.DependencyAnalysis
@@ -438,6 +439,16 @@ public static unsafe void WriteValue(RelocType relocType, void* location, long v
             }
         }
 
+        public static int GetSize(RelocType relocType)
+        {
+            return relocType switch
+            {
+                RelocType.IMAGE_REL_BASED_DIR64 => 8,
+                RelocType.IMAGE_REL_BASED_RELPTR32 => 4,
+                _ => throw new NotSupportedException(),
+            };
+        }
+
         public static unsafe long ReadValue(RelocType relocType, void* location)
         {
             switch (relocType)

src/coreclr/tools/Common/Compiler/DependencyAnalysis/SortableDependencyNode.cs

@@ -42,7 +42,8 @@ protected enum ObjectNodePhase
             /// affect compiler correctness. Today that includes native layout tables.
             /// </summary>
             Ordered,
-            Unordered
+            Unordered,
+            Late,
         }
 
         protected enum ObjectNodeOrder

src/coreclr/tools/Common/Internal/Runtime/DehydratedData.cs

@@ -0,0 +1,118 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using Debug = System.Diagnostics.Debug;
+
+namespace Internal.Runtime
+{
+    /// <summary>
+    /// Provides functionality to encode/decode dehydrated data instruction stream.
+    /// </summary>
+    /// <remarks>
+    /// The instructions use a variable length encoding and are split in two parts:
+    /// the instruction command kind and command data (payload).
+    /// The payload is an integer. If the instruction kind and payload can fit into a single
+    /// byte, the encoding is one byte. Bigger payloads produce bigger instructions.
+    /// </remarks>
+    internal static class DehydratedDataCommand
+    {
+        public const byte Copy = 0x00;
+        public const byte ZeroFill = 0x01;
+        public const byte RelPtr32Reloc = 0x02;
+        public const byte PtrReloc = 0x03;
+
+        private const byte DehydratedDataCommandMask = 0x03;
+        private const int DehydratedDataCommandPayloadShift = 2;
+
+        private const int MaxRawShortPayload = (1 << (8 - DehydratedDataCommandPayloadShift)) - 1;
+        private const int MaxExtraPayloadBytes = 3;
+        private const int MaxShortPayload = MaxRawShortPayload - MaxExtraPayloadBytes;
+
+        public static byte EncodeShort(int command, int commandData)
+        {
+            Debug.Assert((command & DehydratedDataCommandMask) == command);
+            Debug.Assert(commandData <= MaxShortPayload);
+            return (byte)(command | (commandData << DehydratedDataCommandPayloadShift));
+        }
+
+        public static int Encode(int command, int commandData, byte[] buffer)
+        {
+            Debug.Assert((command & DehydratedDataCommandMask) == command);
+            int remainingData = commandData - MaxShortPayload;
+            if (remainingData <= 0)
+            {
+                buffer[0] = EncodeShort(command, commandData);
+                return 1;
+            }
+
+            int numExtraBytes = 0;
+            for (; remainingData != 0; remainingData >>= 8)
+                buffer[++numExtraBytes] = (byte)remainingData;
+            if (numExtraBytes > MaxExtraPayloadBytes)
+                throw new InvalidOperationException(); // decoder can only decode this many extra bytes
+
+            buffer[0] = (byte)(command | ((MaxShortPayload + numExtraBytes) << DehydratedDataCommandPayloadShift));
+            return 1 + numExtraBytes;
+        }
+
+        public static unsafe byte* Decode(byte* pB, out int command, out int payload)
+        {
+            byte b = *pB;
+            command = b & DehydratedDataCommandMask;
+            payload = b >> DehydratedDataCommandPayloadShift;
+            int extraBytes = payload - MaxShortPayload;
+            if (extraBytes > 0)
+            {
+                payload = *++pB;
+                if (extraBytes > 1)
+                {
+                    payload += *++pB << 8;
+                    if (extraBytes > 2)
+                        payload += *++pB << 16;
+                }
+
+                payload += MaxShortPayload;
+            }
+
+            return pB + 1;
+        }
+
+#if false
+        static void Main()
+        {
+            int command, payload;
+
+            byte[] buf = new byte[5];
+            Debug.Assert(Encode(1, 0, buf) == 1);
+            Debug.Assert(buf[0] == 1);
+            Debug.Assert(D(buf, out command, out payload) == 1 && command == 1 && payload == 0);
+            Debug.Assert(Encode(1, 1, buf) == 1);
+            Debug.Assert(buf[0] == (1 | (1 << DehydratedDataCommandPayloadShift)));
+            Debug.Assert(D(buf, out command, out payload) == 1 && command == 1 && payload == 1);
+            Debug.Assert(Encode(1, 60, buf) == 1);
+            Debug.Assert(buf[0] == (1 | (60 << DehydratedDataCommandPayloadShift)));
+            Debug.Assert(D(buf, out command, out payload) == 1 && command == 1 && payload == 60);
+            Debug.Assert(Encode(1, 61, buf) == 2);
+            Debug.Assert(buf[0] == (1 | ((MaxShortPayload + 1) << DehydratedDataCommandPayloadShift)));
+            Debug.Assert(buf[1] == 1);
+            Debug.Assert(D(buf, out command, out payload) == 2 && command == 1 && payload == 61);
+
+            Debug.Assert(Encode(3, 256, buf) == 2);
+            Debug.Assert(D(buf, out command, out payload) == 2 && command == 3 && payload == 256);
+            Debug.Assert(Encode(3, 6500, buf) == 3);
+            Debug.Assert(D(buf, out command, out payload) == 3 && command == 3 && payload == 6500);
+            Debug.Assert(Encode(3, 65000, buf) == 3);
+            Debug.Assert(D(buf, out command, out payload) == 3 && command == 3 && payload == 65000);
+            Debug.Assert(Encode(3, 100000, buf) == 4);
+            Debug.Assert(D(buf, out command, out payload) == 4 && command == 3 && payload == 100000);
+
+            static unsafe int D(byte[] bytes, out int command, out int payload)
+            {
+                fixed (byte* pBytes = bytes)
+                    return (int)(Decode(pBytes, out command, out payload) - pBytes);
+            }
+        }
+#endif
+    }
+}

src/coreclr/tools/Common/Internal/Runtime/ModuleHeaders.cs

@@ -81,7 +81,7 @@ public enum ReadyToRunSectionType
         TypeManagerIndirection = 204,
         EagerCctor = 205,
         FrozenObjectRegion = 206,
-        // 207 is unused - it was used by GCStaticDesc
+        DehydratedData = 207,
         ThreadStaticOffsetRegion = 208,
         // 209 is unused - it was used by ThreadStaticGCDescRegion
         // 210 is unused - it was used by ThreadStaticIndex

src/coreclr/tools/aot/ILCompiler.Compiler.Tests/DependencyGraphTests.cs

@@ -73,7 +73,7 @@ public void TestDependencyGraphInvariants(EcmaMethod method)
                 new FullyBlockedMetadataBlockingPolicy(), new FullyBlockedManifestResourceBlockingPolicy(),
                 null, new NoStackTraceEmissionPolicy(), new NoDynamicInvokeThunkGenerationPolicy(),
                 new ILLink.Shared.TrimAnalysis.FlowAnnotations(Logger.Null, ilProvider, compilerGeneratedState), UsageBasedMetadataGenerationOptions.None,
-                Logger.Null, Array.Empty<KeyValuePair<string, bool>>(), Array.Empty<string>(), Array.Empty<string>(), Array.Empty<string>());
+                default, Logger.Null, Array.Empty<KeyValuePair<string, bool>>(), Array.Empty<string>(), Array.Empty<string>(), Array.Empty<string>());
 
             CompilationBuilder builder = new RyuJitCompilationBuilder(context, compilationGroup)
                 .UseILProvider(ilProvider);

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/AnalysisBasedMetadataManager.cs

@@ -35,7 +35,7 @@ public AnalysisBasedMetadataManager(CompilerTypeSystemContext typeSystemContext)
                 new NoDynamicInvokeThunkGenerationPolicy(), Array.Empty<ModuleDesc>(),
                 Array.Empty<ReflectableEntity<TypeDesc>>(), Array.Empty<ReflectableEntity<MethodDesc>>(),
                 Array.Empty<ReflectableEntity<FieldDesc>>(), Array.Empty<ReflectableCustomAttribute>(),
-                Array.Empty<MetadataType>())
+                Array.Empty<MetadataType>(), default)
         {
         }
 
@@ -51,8 +51,9 @@ public AnalysisBasedMetadataManager(
             IEnumerable<ReflectableEntity<MethodDesc>> reflectableMethods,
             IEnumerable<ReflectableEntity<FieldDesc>> reflectableFields,
             IEnumerable<ReflectableCustomAttribute> reflectableAttributes,
-            IEnumerable<MetadataType> rootedCctorContexts)
-            : base(typeSystemContext, blockingPolicy, resourceBlockingPolicy, logFile, stackTracePolicy, invokeThunkGenerationPolicy)
+            IEnumerable<MetadataType> rootedCctorContexts,
+            MetadataManagerOptions options)
+            : base(typeSystemContext, blockingPolicy, resourceBlockingPolicy, logFile, stackTracePolicy, invokeThunkGenerationPolicy, options)
         {
             _modulesWithMetadata = new List<ModuleDesc>(modulesWithMetadata);
             _typesWithRootedCctorContext = new List<MetadataType>(rootedCctorContexts);

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/CompilationBuilder.Aot.cs

@@ -22,6 +22,7 @@ public partial class CompilationBuilder
         protected bool _methodBodyFolding;
         protected InstructionSetSupport _instructionSetSupport;
         protected SecurityMitigationOptions _mitigationOptions;
+        protected bool _dehydrate;
         protected bool _useDwarf5;
 
         partial void InitializePartial()
@@ -84,6 +85,12 @@ public CompilationBuilder UseSecurityMitigationOptions(SecurityMitigationOptions
             return this;
         }
 
+        public CompilationBuilder UseDehydration(bool dehydrate)
+        {
+            _dehydrate = dehydrate;
+            return this;
+        }
+
         public CompilationBuilder UseMethodBodyFolding(bool enable)
         {
             _methodBodyFolding = enable;

src/coreclr/tools/aot/ILCompiler.Compiler/Compiler/DependencyAnalysis/DehydratableObjectNode.cs

@@ -0,0 +1,29 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace ILCompiler.DependencyAnalysis
+{
+    public abstract class DehydratableObjectNode : ObjectNode
+    {
+        public sealed override ObjectNodeSection GetSection(NodeFactory factory)
+        {
+            return factory.MetadataManager.IsDataDehydrated ? ObjectNodeSection.HydrationTargetSection : GetDehydratedSection(factory);
+        }
+
+        public sealed override ObjectData GetData(NodeFactory factory, bool relocsOnly = false)
+        {
+            ObjectData result = GetDehydratableData(factory, relocsOnly);
+
+            // If we're not generating full data yet, or dehydration is not active,
+            // return the ObjectData as is.
+            if (relocsOnly || !factory.MetadataManager.IsDataDehydrated)
+                return result;
+
+            // Otherwise return the dehydrated data
+            return factory.MetadataManager.PrepareForDehydration(this, result);
+        }
+
+        protected abstract ObjectNodeSection GetDehydratedSection(NodeFactory factory);
+        protected abstract ObjectData GetDehydratableData(NodeFactory factory, bool relocsOnly = false);
+    }
+}