This is the fifth part of the Custom memory allocation series. For your convenience you can find other parts in the table of contents in Part 1 — Allocating object on a stack
In this post we will see how to allocate object on a caller’s stack. We already saw how to allocate object on a stack, but we are unable to access stack’s part from the calling method.
Table of Contents
Introduction
I hope you remember what we were doing in Part 1 of this series. I presented a code for allocating object on a stack using stackalloc and some dirty memory hacks. However, what we really want to do is to invoke something like this:
|
1 |
CustomType myObject = allocator.StackAlloc(() => new CustomType(1, 2, 3)); |
And now we have a problem. Since method’s arguments are passed via stack, the caller is modifying the stack before calling the method. This means that we are unable to allocate object on a stack inside our custom method, because caller will have no option to access this memory when we exit the function.
We could try to inline our function. There is a MethodImplOptions attribute to ask .NET to inline method, however, we have no guarantee that it will be done. JIT might decide that it is too expensive to inline a method and we are out of luck. So we need to modify the caller.
Of course, we could as caller to prepare some memory and give us handle to it:
|
1 2 3 4 |
unsafe{ var memory = stackalloc byte[100]; var myObject = allocator.StackAlloc(() => new CustomType(1, 2, 3), (int)memory); } |
This would work but image how cumbersome it is. We need to add unsafe code everywhere, we need to manually allocate the array, we need to take care of proper casting. Too much hassle.
But imagine the following: the caller simply calls our method and does nothing else. However, during the compile time we modify the code using AOP and manually injects required IL. And this is the way we are going to solve this problem.
Fody to the rescue
Fody is a library for handling manual IL modification. Using this library we can easily add any code to our assemblies, update PDB files, and take care of things which needs to be done but are boring (like logging entering the function). We will write very simple module using Fody to add the code to the calling side.
Let’s begin with the following allocator:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
using CustomMemoryAllocator; using System; using System.Linq.Expressions; namespace CustomStackAllocator { public class BasicStackAllocator : IAllocator { public T Alloc< T>(Expression< Func< T>> allocator) where T : class { // Just a stub method to conform to the IAllocator interface. Should be replaced by Fody or something similar return null; } public T Alloc< T>(Expression< Func< T>> allocator, int address) where T : class { T newObject = (T)AllocationHelper.AllocateAt(address, typeof(T).TypeHandle.Value); GC.SuppressFinalize(newObject); AllocationHelper.InvokeConstructor(newObject, allocator); return newObject; } public void Dispose() { } public void Free< T>(T memory) where T : class { AllocationHelper.InvokeDestructor(memory); } } } |
We have two functions with the same name and different parameters. First one (accepting lambda only) is the function the caller will call. Next, during compilation time we will use Fody to modify the caller to actually call second function (with lambda and integer). The latter function is the one which actually allocates the object.
So we have the following plan:
- We ask the caller to call function accepting lambda only — this way the caller will have no idea how everything works, will have compile time support, and will not need to add unsafe code
- During compile time we find all invocations of our method
- We create byte array on the stack in the invocation place
- Finally, we call second method and pass address of memory
Simple as that. Let’s begin.
Fody’s module
In our main project (the one which allocates memory on a stack) we install Fody from Nuget. Next, we create another project in solution called Weavers with the following class:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
using Mono.Cecil; using Mono.Cecil.Cil; using System; using System.Linq; namespace Weavers { public class ModuleWeaver { public ModuleDefinition ModuleDefinition { get; set; } TypeDefinition stackAllocatorType; MethodDefinition stubMethod; MethodDefinition actualMethod; TypeReference arrayType; TypeDefinition arrayTypeDefinition; public void Execute() { var allTypes = ModuleDefinition.Types.Concat(ModuleDefinition.GetTypeReferences().Select(t => t.Resolve())).ToArray(); stackAllocatorType = allTypes.First(t => t.Name == "BasicStackAllocator"); stubMethod = stackAllocatorType.Methods.First(m => m.Name == "Alloc" && m.Parameters.Count == 1); actualMethod = stackAllocatorType.Methods.First(m => m.Name == "Alloc" && m.Parameters.Count == 2); arrayType = ModuleDefinition.ImportReference(typeof(byte*)); ModuleDefinition.ImportReference(actualMethod); arrayTypeDefinition = arrayType.Resolve(); var methodsToPatch = ModuleDefinition.Types.SelectMany(t => t.Methods.Where(m => m.Body.Instructions.Any(i => i.OpCode == OpCodes.Callvirt && (i.Operand as MethodReference).DeclaringType.Resolve() == stackAllocatorType && (i.Operand as MethodReference).Name == "Alloc" && (i.Operand as MethodReference).Parameters.Count == 1))); foreach (var method in methodsToPatch) { PatchMethod(method); } } } } |
Fody requires our module to have method Execute which does the job. In this method we first get all types and all referenced types from our solution. Next, we look for two methods in our allocator: they have the same name and different number of parameters. Next, we define a type of unsafe array and import it to module. Finally, we examine all methods in module and look far calls to our allocator. Please be aware that the caller now is not able to call our allocator via interface IAllocator, the call must be invoked on the type directly so we can find the line. Next, we patch all methods.
Since we need to generate IL by hand, let’s first write a C# code which we want to have. We start with this:
|
1 2 3 4 5 6 |
using (BasicStackAllocator allocator = new BasicStackAllocator()) { Expression< Func< CustomType>> lambda = () => new CustomType(); var customObject = allocator.Alloc(lambda); /// ... } |
And we want to get the following:
|
1 2 3 4 5 6 7 8 9 10 |
using (BasicStackAllocator allocator = new BasicStackAllocator()) { Expression< Func< CustomType>> lambda = () => new CustomType(); unsafe { var memory = stackalloc byte[100]; var customObject = allocator.Alloc(lambda, (int)memory); /// ... } } |
So we write the latter code, compile it, verify that it is working, and start ILdasm or other decompiler to see the IL. Here it goes:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
; Byte array is at index 2 ; Byte array allocation IL_024b: ldc.i4 100 IL_0250: conv.u IL_0251: ldc.i4.1 IL_0252: mul.ovf.un IL_0253: localloc IL_0255: stloc.2 ; Loading parameters IL_0258: ldloc.0 IL_0259: ldloc.1 IL_0260: ldloc.2 IL_021: conv.i4 ; And here goes method call |
We push the size of the array on the stack, allocate it, and stores pointer to it in variable with index 2. Next, we load parameters to our method (reference to object, lambda, and memory pointer), convert memory pointer to integer, and call the method.
We use the following code:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
for (int i = 0; i < method.Body.Instructions.Count; ++i) { if (method.Body.Instructions[i].OpCode == OpCodes.Callvirt) { var methodTarget = method.Body.Instructions[i].Operand as MethodReference; if (methodTarget != null && methodTarget.DeclaringType.Resolve() == stackAllocatorType && methodTarget.Name == "Alloc" && methodTarget.Parameters.Count == 1) { var existingInstruction = method.Body.Instructions[i]; // Prepare method methodTarget.Parameters.Add(new ParameterDefinition(ModuleDefinition.ImportReference(typeof(int)))); ModuleDefinition.ImportReference(methodTarget); // Define variable var variableDefinition = new VariableDefinition(arrayTypeDefinition); // Remove instruction method.Body.Instructions.RemoveAt(i); // Add variable (byte*) method.Body.Variables.Add(variableDefinition); // Allocate array method.Body.Instructions.Insert(i - 2, Instruction.Create(OpCodes.Ldc_I4_S, (sbyte)100)); method.Body.Instructions.Insert(i - 1, Instruction.Create(OpCodes.Conv_U)); method.Body.Instructions.Insert(i, Instruction.Create(OpCodes.Ldc_I4_1)); method.Body.Instructions.Insert(i + 1, Instruction.Create(OpCodes.Mul_Ovf_Un)); method.Body.Instructions.Insert(i + 2, Instruction.Create(OpCodes.Localloc)); method.Body.Instructions.Insert(i + 3, Instruction.Create(OpCodes.Stloc, variableDefinition)); // Add parameter to method call method.Body.Instructions.Insert(i + 6, Instruction.Create(OpCodes.Ldloc, variableDefinition)); method.Body.Instructions.Insert(i + 7, Instruction.Create(OpCodes.Conv_I4)); //// Modify call target method.Body.Instructions.Insert(i + 8, Instruction.Create(OpCodes.Callvirt, methodTarget)); // Skip added instructions i += 8; } } } |
We iterate through method instructions and look for virtual call to our method. Next, we take method call and add another parameter to it. Next, we remove the existing instruction, declare variable (without name — will be important later), add instructions for array allocations, method to call and finally we call the target.
You might wonder why we are inserting instructions to positions i-2, i-1, i, i+1, i+2, i+3, i+6, i+7, i+8 instead of simply one by one to i, i+1, i+2, i+3, i+4, i+5, i+6, i+7, i+8. .NET considers every C# instruction as more or less transactional — we are not allowed to leave anything on the stack between instructions. Since in our call instruction we have two parameters on the stack already, we cannot insert here code for allocating byte array. We need to allocate this array before we push two parameters on the stack (these are positions i-2, i-1, i, i+1, i+2, i+3), and after that we can push parameters (since we have two instructions for pushing parameters in positions i+4, i+5, we put our instructions from i+6 position). Also please remember, that adding or removing instruction moves other instructions in the method, that is why we modify the instruction line.
We are good to go. We compile the code, check that Fody modified the assembly, start it, and we get unhandled exception. It looks like something is wrong.
Debugging Fody
Since messing with IL is pretty dangerous, we need to take appropriate tools.We get our modified assembly, decompile it, and see the following:
See the type of generated variable — it is uint8 instead of uint8*. It looks like a bug in Fody which you can track here.
So what can we do? Well, we can easily patch the code by hand. We dump the generated IL using ILdasm, edit in using any notepad, add missing star, and compile the code back using ILasm. Because our variable didn’t have a name, we need to fix another error with stloc and ldloc instructions. We compile the code once again and everything works fine.
Summary
We are now able to allocate objects on a stack using simple library function. Fody is pretty good in IL manipulations, so we can use it to easily add other tweaks to code. What’s more, Fody’s modules can be deployed as Nuget packages, so they can be distributed really easily.