This is the fourth 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

We already know how to allocate objects on a stack, how to hijack new operator, how to put objects in any place in memory. Today we are going to write simple trivial implementation of memory allocator. Let’s begin.

Introduction

In theory, we already know how to do that. We know how to allocate a piece of memory and how to put object there. However, for now we were working only with simple objects with no constructors (or with empty constructors). We know how to create object somewhere, but we don’t know how to invoke constructor for it without hijacking new operator.

We will write a simple library capable of doing exactly that. We want the following syntax:

ObjectOfAnyType myObject = allocator.Allocate(() => new ObjectOfAnyType(parameter1, parameter 2, restOfParameters));

So we want to have two things here: first, we want to have a compile time support whether we create object in a correct way (with all parameters correct). Second, we do not want to have any casting. Third, we want to allocate exactly one object — since object creation might be expensive, we do not want to create it on side and copy it somewhere else.

What’s more, our objects need to be casual objects — with constructors, base classes, finalizers, disposable etc. There is no use in memory allocator which is not able to allocate common objects.

Plan

As we can see, we are passing a lambda to allocator method. This takes first requirement: compile time support. Since this lambda is checked by the compiler, we are not able to pass incorrect arguments. What’s more, since this is a lambda, it is not invoked. It only contains all necessary information to be able to execute correctly. We will parse this lambda and extract all informations to create an object.

So our allocator should work as follows:

  • We obtain type of object to create
  • We calculate its size and allocate memory
  • We create an object (by storing its metadata in place)
  • We parse lambda and execute object’s constructor with correct arguments

Everything looks good and we already know how to finish thirst three bullets. Let’s get to the code.

Easier part

We start with the following interface:

using System;
using System.Linq.Expressions;

namespace CustomMemoryAllocator
{
    public interface IAllocator : IDisposable
    {
        T Alloc< T>(Expression< Func< T>> allocator) where T : class;
        void Free< T>(T memory) where T : class;
    }
}

We have three methods: first one is the one which we will use. Second one is there so we will be able to free our object.

Let’s see actual implementation:

using System;
using System.Collections.Generic;
using System.Linq;
using System.Linq.Expressions;
using System.Runtime.InteropServices;

namespace CustomMemoryAllocator
{
    public class BasicHeapAllocator : IAllocator
    {
        private ISet< Tuple< int, int>> _freeChunks = new SortedSet< Tuple< int, int>>();
        private ISet< Tuple< int, int>> _occupiedChunks = new SortedSet< Tuple< int, int>>();
        private IList< byte[]> _chunks = new List< byte[]>();
        public int PageSize { get; private set; }
        private IList< GCHandle> _handles = new List< GCHandle>();
        private bool disposedValue = false;

        public BasicHeapAllocator(int pageSize = 8192)
        {
            PageSize = pageSize;
        }

        public T Alloc< T>(Expression< Func< T>> allocator) where T : class
        {
            T newObject = (T)Allocate(typeof(T).TypeHandle.Value);
            AllocationHelper.InvokeConstructor(newObject, allocator);
            return newObject;
        }

        private object Allocate(IntPtr typeHandle)
        {
            int objectSize = AllocationHelper.GetSize(typeHandle);
            Tuple< int, int> chunk = GetFreeChunk(objectSize);
            // Store chunks ...
            var newObject = AllocationHelper.AllocateAt(chunk.Item1, typeHandle);
            GC.SuppressFinalize(newObject);
            _handles.Add(GCHandle.Alloc(newObject));
            return newObject;
        }

        private Tuple< int, int> AllocateChunk(int bytesToAllocate)
        {
            var memory = new byte[bytesToAllocate];
            _chunks.Add(memory);
            _handles.Add(GCHandle.Alloc(memory));
            return Tuple.Create(AllocationHelper.GetAddress(memory) + 2*sizeof(int), bytesToAllocate);
        }

        public void Free< T>(T memory) where T : class
        {
            var address = AllocationHelper.GetAddress(memory) - 1 * sizeof(int);
            // Update chunks lists ...
            AllocationHelper.InvokeDestructor(memory);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (disposedValue) return;
            disposedValue = true;
            if (!disposing) return;
            
            foreach(var handle in _handles)
            {
                handle.Free();
            }
            for(int i=0;i<_chunks.Count; ++i)
            {
                _chunks[i] = null;
            }
            _freeChunks.Clear();
            _occupiedChunks.Clear();
        }
        
        public void Dispose()
        {
            Dispose(true);
        }
    }
}

Lots of code here. Let's go line by line.

First, we define variables. We store lists of free and occupied chunks as a tuples of integers: first element of a tuple stores the address of the chunk, second element stores the size of the chunk. Next, we store size of a page, a smallest size of allocated chunk we handle. Next, we stores handles to the chunks — we will need to pin them in memory so they are not moved by GC. Finally, flag for IDisposable pattern.

Next goes two important methods. First one dispatches object allocation and invokes constructor using helper class, second one performs actual allocation and performs some accounting — we do not want GC to run destructor of newly created object, we also do not want it to be moved. Rest of the code should be pretty straightforward.

OK, let's see helper classes:

using System;
using System.Linq.Expressions;
using System.Linq;
using System.Runtime.InteropServices;
using System.Reflection;

namespace CustomMemoryAllocator
{
    public static class AllocationHelper
    {
        public static int GetSize< T>() where T : class
        {
            var methodTable = typeof(T).TypeHandle.Value;
            return GetSize(methodTable);
        }

        public static int GetSize(IntPtr typeHandle)
        {
            int objectSize = Marshal.ReadInt32(typeHandle, 4) + 1 * sizeof(int);
            return objectSize;
        }

        public static int GetAddress< T>(T memory)
        {
            unsafe
            {
                TypedReference memoryReference = __makeref(memory);
                var memoryAddress = (*(int*)(*(int*)(&memoryReference)));
                return memoryAddress;
            }
        }

        public static object AllocateAt(int address, IntPtr typeHandle)
        {
            Marshal.WriteInt32((IntPtr)(address), 0);
            Marshal.WriteInt32((IntPtr)(address + 1*sizeof(int)), (int)typeHandle);
            object createdObject = null;
            unsafe
            {
                TypedReference newObjectReference = __makeref(createdObject);
                *(int*)(*(int*)&newObjectReference) = address + 1 * sizeof(int);
            }
            return createdObject;
        }
    }
}

Most of this code is already covered in this series so should be pretty easy. One important thing to notice here is: we do not zero memory for object (we only zero its syncblock). In production code you probably would want to clear the memory before invoking constructor.

OK, easy part is done. Now we need to parse lambda and invoke constructor.

Harder part

We start with the following:

using System;
using System.Linq.Expressions;
using System.Linq;
using System.Runtime.InteropServices;
using System.Reflection;

namespace CustomMemoryAllocator
{
    public static class AllocationHelper
    {
        public static void InvokeConstructor< T>(T target, Expression< Func< T>> allocator) where T : class
        {
            if (allocator.Body.NodeType != ExpressionType.New)
            {
                throw new ArgumentException("Parameter must create an object using new", "allocator");
            }
            var castedExpression = (NewExpression)allocator.Body;
            var constructorInvoker = LambdaHelper.CreateDelegate< T>(castedExpression);
            constructorInvoker.Invoke(target, castedExpression.Arguments.Select(a => LambdaHelper.GetLambda(a).Invoke()).ToArray());
        }
    }
}

First, we check whether passed lambda is in form of new something() If it is not then we abort. Next, we create method for invoking constructor and invoke it.

Before we dive into lambda stuff, we need to consider arguments passing. Since lambda might be very complex (instead of just new MyType(1, 2, reference) we might get something like new MyType(1 * 123 + 15, AbstractProviderFactory.GetObject(ref variable), (() => DateTime.Now)())), we do not want to parse every bit of it. Instead, we are going to invoke the lambda. But since we cannot invoke it as a whole (because it would create an object), we will create a lambda for every single parameter passed to new type() and invoke all of them to get the arguments. Simple as that.

OK, so first let's see LambdaHelper.GetLambda:

public static Func< object> GetLambda(Expression a)
{
   return Expression.Lambda< Func< object>>(Expression.Block(Expression.Convert(a, typeof(object)))).Compile();
}

We get a piece of lambda responsible for a single parameter. It might be something simple (like constant) or something really complex (with asyncs, factories, reflection, etc). We then wrap this lambda in a function returning object. This function contains exactly one additional instruction: casting retrieved parameter to object. Yes — we might be boxing here. But it is actually good to do that.

Let's see an example. Imagine that we create an object in this way:

() => new MyType(AbstractFactory.GetParameter())

In order to obtain the parameter, we first extract part of lambda working on that parameter (namely AbstractFactory.GetParameter()). Next, we create on the fly the following function:

object SomeFunction(){
    return (object)(AbstractFactory.GetParameter());
}

So we have a function which is capable of extracting the argument. Notice that arguments will be evaluated from left to right and they will be evaluated as usual.

Let's get back to constructor invocation:

constructorInvoker.Invoke(target, castedExpression.Arguments.Select(a => LambdaHelper.GetLambda(a).Invoke()).ToArray());

We take all parameters, convert them to functions, compile all of them, and invoke them. Next, we store them in one array of objects and pass to another lambda. So the function we are we want to call should look like this:

// Not C# code!
void ConstructorInvoker< T>(T target, object[] arguments){
    new target(arguments[0], arguments[1], arguments[2], ..., arguments[n]);
}

So we accept target object to invoke constructor and object array with arguments. Next, we need to flatten this array and pass arguments one by one.

The code is as follows:

public static Action< T, object[]> CreateDelegate< T>(NewExpression expression)
        {
            var helperMethod = new DynamicMethod(string.Empty,
                typeof(void),
                new[] { typeof(T), typeof(object[]) },
                typeof(T).Module,
                true);
            var ilGenerator = helperMethod.GetILGenerator();
            ilGenerator.Emit(OpCodes.Ldarg_0);
            for (int i = 0; i < expression.Arguments.Count(); ++i)
            {
                var argumentType = expression.Arguments[i].Type;
                ilGenerator.Emit(OpCodes.Ldarg_1);
                ilGenerator.Emit(OpCodes.Ldc_I4, i);
                ilGenerator.Emit(OpCodes.Ldelem, typeof(object));
                if (argumentType.IsValueType)
                {
                    ilGenerator.Emit(OpCodes.Unbox_Any, argumentType);
                }
            }
            ilGenerator.Emit(OpCodes.Call, expression.Constructor);
            ilGenerator.Emit(OpCodes.Ret);
            var constructorInvoker = (Action< T, object[]>)helperMethod.CreateDelegate(typeof(Action< T, object[]>));
            return constructorInvoker;
        }

First, we create dynamic method accepting two parameters, as described before. Next, we get IL generator and start generating code. We could do the same using reflection as well.

First, we push object reference on the stack. Next, we iterate over all arguments, push on the stack the number of argument (basically, the index in the array), we then load the element from the array and push it on the stack. Finally, we check whether the argument is value type — in that case we unbox the argument.

Next, emit call to constructor and instruction to exit from the function. Finally, we create delegate to the method.

Destructor

We already know how to create an object with lambda, now we need to clean up. Here goes destructor call:

public static void InvokeDestructor< T>(T memory) where T : class
        {
            var toDispose = memory as IDisposable;
            if (toDispose != null) toDispose.Dispose();
            var destructor = memory.GetType().GetMethod("Finalize", BindingFlags.NonPublic | BindingFlags.Instance);
            if (destructor != null)
            {
                destructor.Invoke(memory, null);
            }
        }

We first try to call Dispose from IDisposable interface. Next, we get handler to Finalize method using reflection and invoke it. This method is the destructor.

Summary

Being able to generate code in runtime is really useful. It is not as useful now when we have lambdas and we can generate them on the fly, however, with raw IL code we can do some really nice tricks. In the next part we are going to revisit once again stack allocation.