Friday, August 12, 2011

Writing a Manual Memory Manager in C#

Garbage collection. Aye? or Nay?
As usual, it depends. That is, on which developer you might ask. Some like to have as much control as possible over the way their code executes, while others simply love the fact that they don't have to deal with the "mundane" job of keeping track on their memory allocations.
Since there aren't really any "absolute truths" in anything related to programming, in reality you'd sometimes want to have complete control over your memory management, while at other times you wouldn't really care about it "as long as it gets done".
Since we're mostly discussing .Net and here, we could say that we've got the "as long as it gets done" part covered quite well, by the CLR's garbage collection mechanism. So it's time to see how we could approach implementing a manual memory manager in C#.

What we've eventually would like to have, is an API that would enable us to allocate and deallocate typed memory on demand. Of course C# doesn't natively support the new and delete keywords we so kindly remember from C++, so we'll have come up with our own utility functions to do the job.
Eventually, our code should look something similar to this:

static void Main()
{
    ITypeAllocator mgr = new ManalocManager();

    IFooData obj = mgr.New<IFooData>();

    obj.Bar = 1;
    obj.Car = 2;
    obj.Jar = 3;
    obj.Tar = 4;

    mgr.Delete(obj);
}

Disabling the garbage collector completely is an unreasonable thing to do in a platform such as .Net. Doing so would probably miss the platform's purpose. Anyone who truly wants to have complete control over the execution of its program wouldn't bother using C# anyway (or any other managed language for that matter).
However, while using C#, there might be some times that we'll want to manage our own memory, instead of having the garbage collector doing it for us. And even if not, it's still a subject interesting enough to explore and mostly play with.

In order to demonstrate how we could do achieve manual memory management in C#, lets have a look at the following interface:

public interface IFooData
{
    int Bar { get; set; }
    long Jar { get; set; }
    double Car { get; set; }
    byte Tar { get; set; }
}

The classic method to implement this interface would be to create a class with four members that will match the coresponding properties. However, doing so will result in a 21 bytes structure that will reside in the GC heap (not counting padding and the preceding object header).
Instead, we could allocate the required memory block in the native heap (using AllocHGlobal) and modify out propertie's to access the native memory at the required offsets (e.g. 0 for Bar, 4 for Jar and 12 for Car). Using a Delete method, we could free the native memory block on demand, when we please.

public class FooData : IFooData
{
    private unsafe byte* _native;

    public FooData()
    {
        unsafe
        {
            _native = (byte*)(Marshal.AllocHGlobal(21).ToPointer());
        }
    }

    public int Bar
    {
        get { unsafe { return *(int*)&_native[0]; } }
        set { unsafe { *(int*)&_native[0] = value; } }
    }

    public long Jar
    {
        get { unsafe { return *(long*)&_native[4]; } }
        set { unsafe { *(long*)&_native[4] = value; } }
    }

    public double Car
    {
        get { unsafe { return *(double*)&_native[12]; } }
        set { unsafe { *(double*)&_native[12] = value; } }
    }

    public byte Tar
    {
        get { unsafe { return *(byte*)&_native[20]; } }
        set { unsafe { *(byte*)&_native[20] = value; } }
    }

    public void Delete()
    {
        unsafe
        {
            Marshal.FreeHGlobal(new IntPtr(((void*)(_native))));
            _native = (byte*)(IntPtr.Zero);
        }
    }
}

The problem with such implementation is that it could be very tedious to code and implement. Even for simple structures like IFooData, the resulting implementation could be quite taunting.
Fortunately enough, we can automate the implementation process by adding a code generator that will implemenet our interfaces on the fly, during runtime.
The following interface should loosely describe the capabilities our manual memory manager should support:
 
public interface ITypeAllocator
{
    T New();
    void Delete(T instance);

    void PreGenerate(params Type[] types);
}

The generic parameter T accepts user-defined data representing interfaces such as IFooData.
Once the New method is called, our manager should generate code, compile it and instancate it during runtime. The resulting instance is then returned to the caller for it to be used. Once it finishes using it, and wants to release its memory, it calls the Delete method.
The PreGenerate method's purpose is the optimize the code's generation/compilation process. Once the user pre-generates a type, it won't have to wait on the first call to the New method (much like the process of forcing the JIT compiler to execute on your assemblies).

When it comes to code generation, there are basically two ways to choose from: CodeDOM and Templates. Each one of them has its pros and cons, personaly I tend to prefer the CodeDOM way of doing things. While using it could result in quite verbose code, I believe that its easier to maintain in larger projects than templates.
Unforutantly, .Net's CodeDOM model doesn't support unsafe code, so I had to resort to using some workarounds to represent all of the unsafe code blocks.
This should be a good time to mention the Refly library which wraps around .Net CodeDOM API, making it much simpler and innutative to use.
The demonstrated implementation is very naive and limited regarding the kind of types it is able to generate, though it should illustrate the discussed concept.

public class ManalocManager : ITypeAllocator
{
    // key: userType, value: generatedType
    private Dictionary<Type, Type> m_generatedTypesCache;

    public ManalocManager()
    {
        m_generatedTypesCache = new Dictionary<Type, Type>();
    }

    public void Delete(T instance)
    {
        if (!(instance is IManalocGeneratedType))
            throw new ArgumentException("Attempted to delete an unexpected type");

        IManalocGeneratedType generatedType = (IManalocGeneratedType)instance;
        generatedType.Delete();
    }

    public void PreGenerate(params Type[] types)
    {
        foreach (Type curUserType in types)
            generateAndAddToCache(curUserType);
    }

    public T New()
    {
        Type userType = typeof(T);

        Type generatedType;
        bool alreadyGenerated = m_generatedTypesCache.TryGetValue(userType, out generatedType);
        if (!alreadyGenerated)
            generatedType = generateAndAddToCache(userType);

        object result = Activator.CreateInstance(generatedType);
        return (T)result;
    }

    private Type generateAndAddToCache(Type userType)
    {
        Type generatedType = generateProxy(userType);
        m_generatedTypesCache.Add(userType, generatedType);

        return generatedType;
    }

    private Type generateProxy(Type userType)
    {
        NamespaceDeclaration ns;
        string typeName = createType(userType, out ns);

        string sourceFile = generateCode(ns);
        Assembly compiledAssembly = compile(userType, sourceFile);

        Type compiledType = compiledAssembly.GetType(typeName);

        return compiledType;
    }

    private string createType(Type userType, out NamespaceDeclaration namespaceDec)
    {
        PropertyInfo[] userProperties = userType.GetProperties();

        namespaceDec = new NamespaceDeclaration("Manaloc.AutoGenerated");
        namespaceDec.Imports.Add(userType.Namespace);
        ClassDeclaration classDec = namespaceDec.AddClass(userType.Name + "_Manaloc_AutoGenerated");
        classDec.Interfaces.Add(userType);
        classDec.Interfaces.Add(typeof(IManalocGeneratedType));

        FieldDeclaration nativeMember = classDec.AddField("unsafe byte*", "native");

        addConstructor(userProperties, classDec, nativeMember);
        addDeleteMethod(classDec, nativeMember);
        addProperties(userProperties, classDec, nativeMember);

        string typeName = namespaceDec.Name + "." + classDec.Name;
        return typeName;
    }

    private void addConstructor(PropertyInfo[] userProperties, ClassDeclaration classDec, FieldDeclaration nativeMember)
    {
        int totalSize = sumSize(userProperties);

        ConstructorDeclaration ctor = classDec.AddConstructor();
        ctor.Body.Add(Stm.Snippet("unsafe{"));

        ctor.Body.AddAssign(
            Expr.This.Field(nativeMember),
            Expr.Cast(typeof(byte*), Expr.Type(typeof(Marshal)).Method("AllocHGlobal").Invoke(Expr.Prim(totalSize)).
            Method("ToPointer").Invoke()));

        ctor.Body.Add(Stm.Snippet("}"));
    }

    private void addDeleteMethod(ClassDeclaration classDec, FieldDeclaration nativeMember)
    {
        MethodDeclaration disposeMethod = classDec.AddMethod("Delete");
        disposeMethod.Attributes = MemberAttributes.Final | MemberAttributes.Public;

        disposeMethod.Body.Add(Stm.Snippet("unsafe{"));
        disposeMethod.Body.Add(
            Expr.Type(typeof(Marshal)).Method("FreeHGlobal").Invoke(
            Expr.New(typeof(IntPtr), Expr.Cast("void*", Expr.This.Field(nativeMember)))));
        disposeMethod.Body.AddAssign(Expr.This.Field(nativeMember),
            Expr.Cast("byte*", Expr.Type(typeof(IntPtr)).Field("Zero")));
        disposeMethod.Body.Add(Stm.Snippet("}"));
    }

    private void addProperties(PropertyInfo[] userProperties, ClassDeclaration classDec, FieldDeclaration nativeMember)
    {
        int offset = 0;
        foreach (PropertyInfo curProperty in userProperties)
        {
            Type propType = curProperty.PropertyType;
            int propSize = Marshal.SizeOf(propType);

            PropertyDeclaration propDec = classDec.AddProperty(propType, curProperty.Name);
            propDec.Attributes = MemberAttributes.Final | MemberAttributes.Public;

            if (curProperty.CanRead)
                addGetter(nativeMember, offset, propType, propDec);

            if (curProperty.CanWrite)
                addSetter(nativeMember, offset, propType, propDec);

            offset += propSize;
        }
    }

    private void addSetter(FieldDeclaration nativeMember, int offset, Type propType, PropertyDeclaration propDec)
    {
        propDec.Set.Add(Stm.Snippet("unsafe{"));
        propDec.Set.Add(Stm.Snippet("*(" + propType.Name + "*)&"));
        propDec.Set.AddAssign(Expr.This.Field(nativeMember).Item(offset), Expr.Value);
        propDec.Set.Add(Stm.Snippet("}"));
    }

    private void addGetter(FieldDeclaration nativeMember, int offset, Type propType, PropertyDeclaration propDec)
    {
        propDec.Get.Add(Stm.Snippet("unsafe{"));
        propDec.Get.Add(Stm.Snippet("return *(" + propType.Name + "*)&"));
        propDec.Get.Add(Expr.This.Field(nativeMember).Item(offset));
        propDec.Get.Add(Stm.Snippet("}"));
    }

    private string generateCode(NamespaceDeclaration ns)
    {
        string sourceFile = null;

        const string outDir = "ManalocAutoGenerated";
        if (!Directory.Exists(outDir))
            Directory.CreateDirectory(outDir);

        Refly.CodeDom.CodeGenerator generator = new Refly.CodeDom.CodeGenerator();
        generator.CreateFolders = false;
        generator.FileCreated += (object sender, StringEventArgs args) => { sourceFile = args.Value; };

        generator.GenerateCode(outDir, ns);

        if (sourceFile == null)
            throw new Exception("Faliled to generate source file");

        return sourceFile;
    }

    private Assembly compile(Type userType, string sourceFile)
    {
        CompilerParameters compilerParams = new CompilerParameters();
        compilerParams.CompilerOptions = "/unsafe /optimize";
        compilerParams.ReferencedAssemblies.Add(userType.Assembly.Location);
        CompilerResults result =
            Refly.CodeDom.CodeGenerator.CsProvider.CompileAssemblyFromFile(compilerParams, new string[] { sourceFile });

        Assembly compiledAssembly = result.CompiledAssembly;
        return compiledAssembly;
    }

    private int sumSize(PropertyInfo[] userProperties)
    {
        int size = 0;
        foreach (PropertyInfo curProperty in userProperties)
            size += Marshal.SizeOf(curProperty.PropertyType);

        return size;
    }
}

public interface IManalocGeneratedType
{
    void Delete();
}