This is the third part of the Traits in C# series. For your convenience you can find other parts in the table of contents in Part 1 — Basic implementation with Fody

Last time we implemented method overriding using traits. Today we are going to handle base method calls.

Test program

We will slightly modify our last test program. Here goes the updated code:

using System;
using TraitIntroducer;


namespace TraitsDemo
{
    public interface IA
    {
    }

    public interface IB
    {
    }

    public interface IC
    {
    }

    [TraitFor(typeof(IA))]
    public static class IA_Implementation
    {
        public static void Print(this IA instance)
        {
            Console.WriteLine("I'm IA");
            instance.Base();
        }
    }

    [TraitFor(typeof(IB))]
    public static class IB_Implementation
    {
        public static void Print(this IB instance)
        {
            Console.WriteLine("I'm IB");
            instance.Base();
        }
    }

    [TraitFor(typeof(IC))]
    public static class IC_Implementation
    {
        public static void Print(this IC instance)
        {
            Console.WriteLine("I'm IC");
            instance.Base();
        }
    }

    [TraitFor(typeof(IC))]
    public static class IC_Implementation2
    {
        public static void Print(this IC instance)
        {
            Console.WriteLine("I'm IC2");
            instance.Base();
        }
    }

    public class A : IA
    {
        public virtual void Print()
        {
            Console.WriteLine("I'm A");
            this.Base();
        }
    }

    public class B : A
    {
    }

    public class C : B
    {
        public override void Print()
        {
            Console.WriteLine("I'm C");
            this.Base();
        }
    }

    public class D : C, IB, IC
    {
    }

    class Program
    {
        static void Main(string[] args)
        {
            IA a = new D();
            a.Print();
			// Should print:
            // I'm IC2
            // I'm IC
            // I'm IB
            // I'm C
            // I'm A
            // I'm IA
        }
    }
}

In ordinary C# we would like to call base methods using base.Print. However, we cannot do that since we don’t have methods in interfaces, they are introduced by Fody during compile time. So what can we do? We can add stub method which we will replace with Fody. Stub method looks like this:

namespace TraitIntroducer
{
    public static class TratExtensions
    {
        public static void Base(this object instance)
        {
        }
    }
}

We extend System.Object which means that we can call this method with any object we like. Instead of calling base.Print, we call this.Base(). In Fody’s weaver we will take care of replacing call to Base with call to concrete method. However, if we do not replace the call, nothing wrong happens — extension method will just do nothing, and JIT should remove the call anyway.

Let’s implement logic.

Weaver

In fact, not much changes. First, we need to find our method:

public void Execute(){
	// ...
	BaseMethodExtension = AllTypes.First(t => t.FullName == typeof(TratExtensions).FullName).Methods.First(m => m.Name == "Base");
	// ...
}
			
private bool IsCallToOurBase(Instruction instruction, string name)
{
        return instruction.OpCode == OpCodes.Call && (instruction.Operand as MethodReference)?.FullName == name;
}

Next, we need to fix logic modifying classes:

private void FixClass(TypeDefinition type)
{
	var hierarchy = GetTypeHierarchy(type).Reverse().Skip(1).TakeWhile(t => t.IsInterface).Reverse().ToArray();
	var upperHierarchy = GetTypeHierarchy(type).Reverse().Skip(1).SkipWhile(t => t.IsInterface).Where(t => t.IsClass).ToArray();
	var introducedMethods = new Dictionary< string, Tuple< MethodDefinition, MethodDefinition>>();
	var lastInjectedMethod = new Dictionary< string, Tuple< MethodDefinition, MethodDefinition>>();
	var methodsToFix = new Dictionary< string, MethodDefinition>();
	foreach (var implementedInterface in hierarchy)
	{
		var extenders = GetExtenders(implementedInterface);
		foreach (var extender in extenders)
		{
			foreach (var method in extender.Methods)
			{
				Tuple< MethodDefinition, MethodDefinition> baseMethodPair;
				lastInjectedMethod.TryGetValue(method.Name, out baseMethodPair);

				FixUpperHierarchy(upperHierarchy, method);
				var baseMethod = baseMethodPair?.Item2 ?? GetMethodByNameFromUpperHierarchy(upperHierarchy, method);
				
				var injected = InjectExtensionMethodToInheritor(type, extender, method, baseMethod);
				var matchingMethod = type.Methods.FirstOrDefault(m => m.Name == method.Name);
				lastInjectedMethod[method.Name] = Tuple.Create(method, injected);
				if (matchingMethod != null)
				{
					matchingMethod.Attributes |= MethodAttributes.Virtual;
					matchingMethod.Attributes &= ~MethodAttributes.NewSlot;
					methodsToFix[matchingMethod.Name] = method;
					FixSingleMethodCallToBase(matchingMethod, injected, baseMethod);
				}
				else
				{
					introducedMethods[method.Name] = Tuple.Create(method, injected);
				}
			}
		}
	}

	foreach (var introducedMethod in introducedMethods.Values)
	{
		InjectVirtualMethodToInheritor(type, introducedMethod.Item1, introducedMethod.Item2);
	}
	foreach (var methodToFix in methodsToFix)
	{
		FixUpperHierarchy(upperHierarchy, methodToFix.Value);
	}
}

We basically have three possibilities when modifying class:

• We don’t touch the class — because it doesn’t implement any interesting interfaces — in this case we need to fix base calls anyway since upper classes might implement some interface
• We only add extension methods — because method with matching name already exists — we need to fix matching method since it might contain calls to base methods, we also need to fix extension methods
• We add everything (extension methods + virtual method) — in that case we only need to fix extension methods

OK, so what’s new? We add a dictionary for existing methods we need to fix. Before injecting a method, we look for matching method from base classes or take already injected method with the same name. We inject extension to inheritor (will come back to this in a second). After introducing all methods, we fix hierarchy in the following way:

private MethodDefinition GetMethodByNameFromUpperHierarchy(TypeDefinition[] upperHierarchy, MethodDefinition method)
{
	return upperHierarchy.Select(t => t.Methods.FirstOrDefault(m => m.Name == method.Name)).FirstOrDefault(m => m != null);
}

private void FixUpperHierarchy(TypeDefinition[] hierarchy, MethodDefinition method)
{
	if (!hierarchy.Any())
	{
		return;
	}

	var typeToFix = hierarchy.First();
	var restHierarchy = hierarchy.Skip(1).ToArray();

	var matchingMethod = typeToFix.Methods.FirstOrDefault(m => m.Name == method.Name);
	var upperMethod = GetMethodByNameFromUpperHierarchy(restHierarchy, method);
	FixSingleMethodCallToBase(matchingMethod, upperMethod);

	FixUpperHierarchy(restHierarchy, method);
}

private void FixSingleMethodCallToBase(MethodDefinition existingMethod, MethodDefinition methodToCall, MethodDefinition baseMethod = null)
{
	if (existingMethod == null || methodToCall == null)
	{
		return;
	}

	foreach (var instruction in existingMethod
		.Body.Instructions
		.Where(i => 
			IsCallToOurBase(i, BaseMethodExtension.FullName) 
			|| (baseMethod != null && IsCallToOurBase(i, baseMethod.FullName))
		)
	)
	{
		instruction.Operand = methodToCall.GetElementMethod();
	}
}

We examine base classes, try to find matching base methods, and replace calls. Please notice, that if there is any call to base method in one of base classes, it means, that this class does not implement interesting interfaces, and its base call should be redirected to virtual method in one of base classes. We do this recursively in order to fix all methods. We might check methods multiple times, but we will modify call to a method once for every new injected implementation.

Finally, we need to inject extension methods in slightly different way:

private MethodDefinition InjectExtensionMethodToInheritor(TypeDefinition inheritor, TypeDefinition extender, MethodDefinition method, MethodDefinition baseMethod)
{
	var newMethod = new MethodDefinition(GetInheritorMethodName(extender, method), MethodAttributes.Public | MethodAttributes.Virtual | MethodAttributes.HideBySig | MethodAttributes.NewSlot, method.ReturnType);
	foreach (var instruction in method.Body.Instructions)
	{
		if (baseMethod != null && IsCallToOurBase(instruction, BaseMethodExtension.FullName))
		{
			newMethod.Body.Instructions.Add(Instruction.Create(OpCodes.Call, baseMethod));
		}
		else
		{
			newMethod.Body.Instructions.Add(instruction);
		}
	}
	inheritor.Methods.Add(newMethod);

	return newMethod;
}

And we are done.

Results

As usual, below are decompiled codes. Since only implementation were changed, we do not examine interfaces this time. Test program:

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.Program
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 294DB379-D664-4EBC-818F-5EBC7551B445
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

namespace TraitsDemo
{
  internal class Program
  {
    private static void Main(string[] args)
    {
      IA_Implementation.Print(new D());
    }
  }
}

Extension classes:

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IA_Implementation
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 294DB379-D664-4EBC-818F-5EBC7551B445
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IA))]
  public static class IA_Implementation
  {
    public static void Print(this IA instance)
    {
      instance.Print();
    }
  }
}


// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IB_Implementation
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 294DB379-D664-4EBC-818F-5EBC7551B445
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IB))]
  public static class IB_Implementation
  {
    public static void Print(this IB instance)
    {
      instance.Print();
    }
  }
}


// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IC_Implementation
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 294DB379-D664-4EBC-818F-5EBC7551B445
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IC))]
  public static class IC_Implementation
  {
    public static void Print(this IC instance)
    {
      instance.Print();
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IC_Implementation2
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 479D6B08-31BC-46C3-807F-322E44C5DECA
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IC))]
  public static class IC_Implementation2
  {
    public static void Print(this IC instance)
    {
      instance.Print();
    }
  }
}

Concrete classes:

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.A
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 479D6B08-31BC-46C3-807F-322E44C5DECA
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using System;
using TraitIntroducer;

namespace TraitsDemo
{
  public class A : IA
  {
    public virtual void Print()
    {
      Console.WriteLine("I'm A");
      this.Print_IA_Implementation();
    }

    public virtual void Print_IA_Implementation()
    {
      Console.WriteLine("I'm IA");
      this.Base();
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.B
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 479D6B08-31BC-46C3-807F-322E44C5DECA
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

namespace TraitsDemo
{
  public class B : A
  {
  }
}


// Decompiled with JetBrains decompiler
// Type: TraitsDemo.C
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 479D6B08-31BC-46C3-807F-322E44C5DECA
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using System;

namespace TraitsDemo
{
  public class C : B
  {
    public override void Print()
    {
      Console.WriteLine("I'm C");
      base.Print();
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.D
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 479D6B08-31BC-46C3-807F-322E44C5DECA
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using System;

namespace TraitsDemo
{
  public class D : C, IB, IC
  {
    public virtual void Print_IB_Implementation()
    {
      Console.WriteLine("I'm IB");
      base.Print();
    }

    public virtual void Print_IC_Implementation()
    {
      Console.WriteLine("I'm IC");
      this.Print_IB_Implementation();
    }

    public virtual void Print_IC_Implementation2()
    {
      Console.WriteLine("I'm IC2");
      this.Print_IC_Implementation();
    }

    public override void Print()
    {
      this.Print_IC_Implementation2();
    }
  }
}

As we can see, everything is correct.

Summary

We have implemented very basic logic for introducing traits to C# codebase. We are now ready to use this feature to stack implementations in pretty nice way. Of course, there is much more to do in order to use it in the production environment.

First, we onlly support void methods without parameters. Handling different methods is not difficult, but it takes time. We could use lambdas to have compile time support, but verifying lambdas in Fody takes even more time.

Next, we only support public virtual methods. It is not difficult to modify behaviour using attributes, but once again — it takes time.

Third, we don’t report errors or check existing code — we blindly substitute methods without verifying whether everything makes sense.

These are of course cases which we should take care before deploying this code to production. But since we do it only for fun, we perform little to no error checking.

You can find whole code at Github

This is the second part of the Traits in C# series. For your convenience you can find other parts in the table of contents in Part 1 — Basic implementation with Fody

Last time we saw simple implementation of traits in C#. Today we are going to handle various overrides. Let’s go.

Test code

We want to write the following application:

using System;
using TraitIntroducer;

namespace TraitsDemo
{
    public interface IA
    {
    }

    public interface IB
    {
    }

    public interface IC
    {
    }

    [TraitFor(typeof(IA))]
    public static class IA_Implementation
    {
        public static void Print(this IA instance)
        {
            Console.WriteLine("I'm IA");
        }
    }

    [TraitFor(typeof(IB))]
    public static class IB_Implementation
    {
        public static void Print(this IB instance)
        {
            Console.WriteLine("I'm IB");
        }
    }

    [TraitFor(typeof(IC))]
    public static class IC_Implementation
    {
        public static void Print(this IC instance)
        {
            Console.WriteLine("I'm IC");
        }
    }

    [TraitFor(typeof(IC))]
    public static class IC_Implementation2
    {
        public static void Print(this IC instance)
        {
            Console.WriteLine("I'm IC2");
        }
    }

    public class A : IA
    {
        public virtual void Print()
        {
            Console.WriteLine("I'm A");
        }
    }

    public class B : A
    {
    }

    public class C : B
    {
        public override void Print()
        {
            Console.WriteLine("I'm C");
        }
    }

    public class D : C, IB, IC
    {
    }

    class Program
    {
        static void Main(string[] args)
        {
            IA a = new D();
            a.Print(); // Should print "I'm IC2"
        }
    }
}

What’s exactly going on here? We have three interfaces (they do not inherit from each other). We also have a simple class hierarchy: first class contains virtual method Print, and implements first interface. Second clas inherits from first class and nothing else. Third class inherits from second class and overrides method. Fourth class inherits from third class and implements two additional interfaces.

We create instance of last class and store its reference with first interface as type. Next, we call method. If we do not wire up traits logic, it should print “I’m IA”, since it should be bound during compilation. However, our hierarchy should look like this:
IA_Implementation.Print -> A.Print -> C.Print -> IB_Implementation.Print -> IC_Implementation.Print -> IC_Implementation2.Print
so we would like to see “I’m IC2” printed out.And this is the task we are going to implemented today: selecting correct override. Next time we will see how to implement base.Print in order to be able to stack traits.

Implementation

Since we want to select overrides, it might be helpful to be able to extract whole inheritance hierarchy. Let’s start with the following code:

private IEnumerable< TypeDefinition> GetTypeHierarchy(TypeDefinition type)
{
	if (type == null)
	{
		return Enumerable.Empty< TypeDefinition>();
	}

	return GetTypeHierarchy(type.BaseType as TypeDefinition)
		.Concat(type.Interfaces.OfType< TypeDefinition>().SelectMany(GetTypeHierarchy))
		.Concat(new [] {type})
		.GroupBy(t => t.FullName)
		.Select(x => x.First());
}

If passed type is null, we return empty collection. In other case we get base type’s hierarchy, add interfaces implemented by us, add ourselves, and finally remove duplicates. So for class D in demo program we would like to have the following class hierarchy:

IA -> A -> B -> C -> IB -> IC -> D

We omit System.Object here.

With this tool in hand we can implement weaver:

public void Execute()
{
	AllTypes = ModuleDefinition.Types.Concat(ModuleDefinition.GetTypeReferences().Select(t => t.Resolve())).ToList();
	TraitForAttributeTypeDefinition = AllTypes.FirstOrDefault(t => t.FullName == typeof(TraitForAttribute).FullName);
	if (TraitForAttributeTypeDefinition == null) return;

	var orderedTypes = AllTypes.OrderBy(type => GetTypeHierarchy(type).Count());

	foreach (var type in orderedTypes.Where(t => t.IsClass))
	{
		FixClass(type);
	}
	foreach (var type in orderedTypes.Where(t => t.IsInterface))
	{
		FixInterface(type);
	}
}

We first get all types from solution, and type for our attribute. Next, we sort types by their hierarchy — we want to modify types in top-down manner so this ordering should be good for us. Next, we fix all classes first, and after that we fix interfaces.

In order to fix anything, we need to be able to extract extension classes for specified interface:

private IEnumerable< TypeDefinition> GetExtenders(TypeDefinition typeToExtend)
{
	return AllTypes
		.Where(type => type.CustomAttributes.Any(attribute => attribute.AttributeType.FullName == TraitForAttributeTypeDefinition.FullName))
		.Where(type =>
		{
			var extendedInterafaceType = type.CustomAttributes.First(attribute => attribute.AttributeType.FullName == TraitForAttributeTypeDefinition.FullName).ConstructorArguments.First().Value as TypeDefinition;
			var extendedInterfaceTypeDefinition = AllTypes.First(t => t.FullName == extendedInterafaceType.FullName);
			return extendedInterfaceTypeDefinition.FullName == typeToExtend.FullName;
		});
}

We iterate over all types, find types implementing interfaces, and check for interface types. Since we are allowed to have multiple extension classes for single interface, we need to return collection.

Now we can fix classes. For each class we need to add methods from extension classes implementing interfaces implemented by this specific class (and not by some ancestor), and add virtual method for each method from interface:

private void FixClass(TypeDefinition type)
{
	var hierarchy = GetTypeHierarchy(type).Reverse().Skip(1).TakeWhile(t => t.IsInterface).Reverse().ToArray();
	var introducedMethods = new Dictionary< string, Tuple< MethodDefinition, MethodDefinition>>();
	foreach (var implementedInterface in hierarchy)
	{
		var extenders = GetExtenders(implementedInterface);
		foreach (var extender in extenders)
		{
			foreach (var method in extender.Methods)
			{
				var injected = InjectExtensionMethodToInheritor(type, extender, method);
				var matchingMethod = type.Methods.FirstOrDefault(m => m.Name == method.Name);
				if (matchingMethod != null)
				{
					matchingMethod.Attributes |= MethodAttributes.Virtual;
					matchingMethod.Attributes &= ~MethodAttributes.NewSlot;
				}
				else
				{
					introducedMethods[method.Name] = Tuple.Create(method, injected);
				}
			}
		}
	}

	foreach (var introducedMethod in introducedMethods.Values)
	{
		InjectVirtualMethodToInheritor(type, introducedMethod.Item1, introducedMethod.Item2);
	}
}

Looks like black magic, so let’s go step by step. We first get hierarchy for modified type. Since we want to introduce only methods from interfaces implemented by our class (and not by ancestor), from hierarchy we take only last interfaces.

Next, we create dictionary for introduced methods — we want to call method introduced at the end, so we need to save its instance for later.

We iterate over all interfaces, over all extension classes, and over all methods. For each such a method we inject it. Injected method will have name in the form {methodName}_{extensionClasName}. So for method Print of interface IC implemented in IC_Implementation2 we will name the method Print_IC_Implementation2.

Next, we check if we have method with matching name already in the class. So in this case we check if we have Print in our class implemented. If it is the case, we fix its attributes — in order to make method virtual. If we don’t have such a method, we save instance of injected method from interface in order to be able to add it later.

Finally, when three nested loops are done, we iterate over dictionary, and add all missing virtual methods.

These are methods injecting code for concrete classes:

private string GetInheritorMethodName(TypeDefinition extender, MethodDefinition method)
{
	return $"{method.Name}_{extender.Name}";
}

private MethodDefinition InjectExtensionMethodToInheritor(TypeDefinition inheritor, TypeDefinition extender, MethodDefinition method)
{
	var newMethod = new MethodDefinition(GetInheritorMethodName(extender, method), MethodAttributes.Public | MethodAttributes.Virtual | MethodAttributes.HideBySig | MethodAttributes.NewSlot, method.ReturnType);
	foreach (var instruction in method.Body.Instructions)
	{
		newMethod.Body.Instructions.Add(instruction);
	}
	inheritor.Methods.Add(newMethod);

	return newMethod;
}

private MethodDefinition InjectVirtualMethodToInheritor(TypeDefinition inheritor, MethodDefinition method, MethodDefinition methodToCall)
{
	var newMethod = new MethodDefinition(method.Name, MethodAttributes.Public | MethodAttributes.Virtual | MethodAttributes.HideBySig, method.ReturnType);
	newMethod.Body.Instructions.Add(Instruction.Create(OpCodes.Ldarg_0));
	newMethod.Body.Instructions.Add(Instruction.Create(OpCodes.Callvirt, methodToCall));
	newMethod.Body.Instructions.Add(Instruction.Create(OpCodes.Ret));
	inheritor.Methods.Add(newMethod);

	return newMethod;
}

Extension method is introduced as is, virtual method calls most recent injected method.

Concrete classes are done. We need to fix interfaces:

private void FixInterface(TypeDefinition type)
{
	var hierarchy = GetTypeHierarchy(type).Reverse().Skip(1).ToArray();
	var extenders = GetExtenders(type);
	foreach (var extender in extenders)
	{
		foreach (var method in extender.Methods)
		{
			var existingMethod = hierarchy.Select(t => t.Methods.FirstOrDefault(m => m.Name == method.Name)).FirstOrDefault(m => m != null) ?? type.Methods.FirstOrDefault(m => m.Name == method.Name);
			if (existingMethod == null)
			{
				InjectMethodToInterface(type, method);
			}
			else
			{
				FixTraitMethod(method, existingMethod);
			}
		}
	}
}

We first calculate hierarchy, and find all extension classes. Next, for each method we first check if any of our ancestor contains method with the same name, or if we already have this method. If it is not the case, we inject method:

private void InjectMethodToInterface(TypeDefinition extendedInterface, MethodDefinition method)
{
	var newMethod = new MethodDefinition(method.Name, MethodAttributes.Abstract | MethodAttributes.Virtual | MethodAttributes.HideBySig | MethodAttributes.NewSlot, method.ReturnType);
	extendedInterface.Methods.Add(newMethod);
	FixTraitMethod(method, newMethod);
}

No magic here, we do the same, as in the last part. We also need to fix method in extension class:

private void FixTraitMethod(MethodDefinition method, MethodDefinition methodToCall)
{
	method.Body.Instructions.Clear();
	method.Body.Instructions.Add(Instruction.Create(OpCodes.Ldarg_0));
	method.Body.Instructions.Add(Instruction.Create(OpCodes.Callvirt, methodToCall));
	method.Body.Instructions.Add(Instruction.Create(OpCodes.Ret));
}

And we are done.

Result

Let’s compile the code and see exactly the result. These are interfaces:

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IA
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 7CA9337F-2598-4715-B2DF-60EC6ECB1D59
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

namespace TraitsDemo
{
  public interface IA
  {
    void Print();
  }
}


// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IB
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 7CA9337F-2598-4715-B2DF-60EC6ECB1D59
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

namespace TraitsDemo
{
  public interface IB
  {
    void Print();
  }
}


// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IC
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 55F275F7-2F8B-40AB-B38A-3F6A73C027AF
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

namespace TraitsDemo
{
  public interface IC
  {
    void Print();
  }
}

Classes:

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.A
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 7CA9337F-2598-4715-B2DF-60EC6ECB1D59
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using System;
using TraitIntroducer;

namespace TraitsDemo
{
  public class A : IA
  {
    public virtual void Print()
    {
      Console.WriteLine("I'm A");
    }

    public virtual void Print_IA_Implementation()
    {
      Console.WriteLine("I'm IA");
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.B
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 7CA9337F-2598-4715-B2DF-60EC6ECB1D59
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

namespace TraitsDemo
{
  public class B : A
  {
  }
}


// Decompiled with JetBrains decompiler
// Type: TraitsDemo.C
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 7CA9337F-2598-4715-B2DF-60EC6ECB1D59
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using System;
using TraitIntroducer;

namespace TraitsDemo
{
  public class C : B
  {
    public override void Print()
    {
      Console.WriteLine("I'm C");
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.D
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 9C15A304-C4AC-4E6C-991B-67F70B172CB3
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using System;
using TraitIntroducer;

namespace TraitsDemo
{
  public class D : C, IB, IC
  {
    public virtual void Print_IB_Implementation()
    {
      Console.WriteLine("I'm IB");
    }

    public virtual void Print_IC_Implementation()
    {
      Console.WriteLine("I'm IC");
    }

    public virtual void Print_IC_Implementation2()
    {
      Console.WriteLine("I'm IC2");
    }

    public override void Print()
    {
      this.Print_IC_Implementation2();
    }
  }
}

Test program:

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.Program
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 9C15A304-C4AC-4E6C-991B-67F70B172CB3
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

namespace TraitsDemo
{
  internal class Program
  {
    private static void Main(string[] args)
    {
      IA_Implementation.Print(new D());
    }
  }
}

And extension classes:

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IA_Implementation
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 9C15A304-C4AC-4E6C-991B-67F70B172CB3
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IA))]
  public static class IA_Implementation
  {
    public static void Print(this IA instance)
    {
      instance.Print();
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IB_Implementation
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 9C15A304-C4AC-4E6C-991B-67F70B172CB3
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IB))]
  public static class IB_Implementation
  {
    public static void Print(this IB instance)
    {
      instance.Print();
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IC_Implementation
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 9C15A304-C4AC-4E6C-991B-67F70B172CB3
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IC))]
  public static class IC_Implementation
  {
    public static void Print(this IC instance)
    {
      instance.Print();
    }
  }
}

// Decompiled with JetBrains decompiler
// Type: TraitsDemo.IC_Implementation2
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 9C15A304-C4AC-4E6C-991B-67F70B172CB3
// Assembly location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe

using System;
using TraitIntroducer;

namespace TraitsDemo
{
  [TraitFor(typeof (IC))]
  public static class IC_Implementation2
  {
    public static void Print(this IC instance)
    {
      instance.Print();
    }
  }
}

So we create instance of class D, call method IA_Implementation.Print which in turn calls instance.Print(), which goes down to D.Print, which finally calls this.Print_IC_Implementation2(), and this prints “I’m IC2”. We are done!

Summary

We saw how to inject correct methods and bind them during compile time. In next part we are going to implement traits stacking, which will allow us to create decorators in really nice way, just like we can do in Scala.
You can find whole code in this gist.
Bonus chatter: what happens if in interface we define two methods void Print()? Will the .NET load type correctly?
Bonus chatter 2: how does base.Print works internally? Does it perform static or dynamic dispatch? What IL instruction does it use?

This is the first part of the Traits in C# series. For your convenience you can find other parts using the links below :
Part 1 — Basic implementation with Fody
Part 2 — Overrides
Part 3 — Stacking traits

If you are interested in the topic see the
article from „Implementacja traitów w języku C# przy użyciu biblioteki Fody” in Programista nr 55

Hi! This post starts new series about implementing traits in C#. Let’s go.

What is that?

Wikipedia says that “trait is a concept used in object-oriented programming, which represents a set of methods that can be used to extend the functionality of a class”. It allows us to reuse existing functions in new types. Trait is something between class and interface — single class can have multiple traits (“inherit” from them), so this is similar to interfaces (since we can implement multiple interfaces in C# class), and each trait can have multiple methods with implementation (just like method in a class). However, traits are stateless — they cannot provide variables, only methods.

As a side note — stating that C# doesn’t support multiple inheritance is incorrect. We cannot inherit from multiple classes, however, we can implement (inherit from) multiple interfaces. It is a good moment to distinct between interface inheritance and implementation inheritance. The former means that we are allowed to gather interface (method signatures) from different types, the latter indicates that we can reuse actual implementation. There are languages supporting inheritance from multiple classes — C++ is an epitomy here. It is also worth noting that there is another important concept called mix—in. Actually, Scala’s traits are in fact mix-ins.

Now you might ask — why do I ever need this kind of a thing? Well, it’s difficult question. You might very well never need to use the concept, but there are situations when it would be very helpful to be able to define trait. Imagine the case: we have IEnumerable with method Select. We have also IQueryable which has the same method (to be precise: method with the same signature). But if you try to do something like this:

IEnumerable< Foo> foos = SomethingReturningIQueryableOfFoos();
foos.Select(f => f.Something);

you will actually call Select from IEnumerable instead of IQueryable, even though foos is the latter. This is because of early binding — Select is just an extension method and it is bound during compilation. However, let’s take this Java code:

import java.util.*;
import java.lang.*;
import java.io.*;
 
interface Interface1 {
   default void print(){
      System.out.println("I am Interface1");
   }
}
 
interface Interface2 extends Interface1 {
   default void print(){
      System.out.println("I am Interface2");
   }
}
 
class Class1 implements Interface2 {
 
}
 
class Ideone
{
	public static void main (String[] args) throws java.lang.Exception
	{
		Interface1 instance = new Class1();
		instance.print();
	}
}

This will actually print “I am Interface2”, even though instance variable has compilation type Interface1. If you try to reimplement this snippet in C# using extension methods, you will get different results.

As we can see, trait is pretty powerful concept. Unfortunately, C# doesn’t support them. In this series we are going to impement very basic Fody module supporting traits.

Usage

Let’s write down what we want to achieve:

• We would like to be able to define a trait
• We would like to have intellisense — we still want to have compiler support when writing code
• We want traits to be polymorphic — they need to support late binding
• For sake of simplicity we won’t bother with fancy functions — we will just support parameterless void ones
• As usual, we want the process to be clean and reliable

We need to have methods’ implementations somewhere. Since we want to have compiler’s support, we will use extension methods. Basically, our code will look like this:

using System;
using TraitIntroducer;

namespace TraitsDemo
{
    public interface IA
    {
    }

    [TraitFor(typeof(IA))]
    public static class IA_Implementation
    {
        public static void Print( this IA instance)
        {
            Console.WriteLine( "I'm IA");
        }
    }

    public class A : IA
    {       
    }

    public class B : A
    {
        public virtual void Print()
        {
            Console.WriteLine( "I'm B");
        }
    }

    class Program
    {
        static void Main( string[] args)
        {
            IA ia = new B();
            ia.Print();
        }
    }
}

If you compile this code as-is, it will print “I’m IA”. Why? Because ia variable is of type IA during compilation, and ia.Print() is bound to static extension method. However, we would like this to be transformed to the following:

using System;
using TraitIntroducer;

namespace TraitsDemo
{
    public interface IA
    {
        void Print();
    }

    [TraitFor(typeof(IA))]
    public static class IA_Implementation
    {
        public static void Print( this IA instance)
        {
            instance.Print();
        }
    }

    public class A : IA
    {    
        public virtual void Print( this IA instance)
        {
            Console.WriteLine( "I'm IA");
        }   
    }

    public class B : A
    {
        public virtual override void Print()
        {
            Console.WriteLine( "I'm B");
        }
    }

    class Program
    {
        static void Main( string[] args)
        {
            IA ia = new B();
            ia.Print();
        }
    }
}

So we introduce method Print to interface IA, we move original implementation from extension method to newly created method in class A, and we modify existing extension method to call method from interface. What’s more, since ia is actually of type B which overrides method Print, we should get different result. This time (thanks to late binding) we expect to get “I’m B”.

Implementation

As we said, we are going to utilize Fody to modify the code. We need to mark classes providing actual implementation in order to be able to find them in weaver. So we start with the following attribute:

using System;

namespace TraitIntroducer
{
    public class TraitForAttribute : Attribute
    {
        public Type InterfaceType { get; private set; }

        public TraitForAttribute(Type interfaceType)
        {
            InterfaceType = interfaceType;
        }
    }
}

This attribute accepts one parameter — type of interface which we want to boost. Using this attribute we will write the code in the following way:

• We define empty interface which we will fill up with methods using Fody
• Methods which we would like to be implemented in interface we will implement as extension methods — thanks to this we will have intellisense
• We will handle rest using Fody

Let’s dig into Fody’s module. We start with the following:

public void Execute()
{
	var allTypes = ModuleDefinition.Types.Concat(ModuleDefinition.GetTypeReferences().Select(t => t.Resolve())).ToArray();
	var traitForAttributeTypeDefinition = allTypes.FirstOrDefault(t => t.FullName == typeof(TraitForAttribute).FullName);
	if (traitForAttributeTypeDefinition == null) return;
	IDictionary< TypeDefinition, IList< MethodDefinition>> introducedMethods = GetIntroducedMethods(allTypes, traitForAttributeTypeDefinition);
	foreach (var tuple in introducedMethods)
	{
		ExtendInterface(tuple.Key, tuple.Value);
		ExtendInheritors(tuple.Key, tuple.Value, allTypes);
		FixTraitMethods(tuple.Key, tuple.Value);
	}
}

First, we get all types in order to examine them. Next, we look for attribute’s metadata. If we can’t find it, it probably means that it is not used, so we are done. If it is used, we group all introduced methods (extension methods) by type of their target interface. For each such method we first extend interface (which means copying extension method signature to the interface), we extend inheritors (which means that we copy extension methods to classes), and fix trait methods (which means that we modify extension methods to call methods from interfaces).

Let’s see how we extract methods to introduce to interfaces:

private IDictionary< TypeDefinition, IList< MethodDefinition>> GetIntroducedMethods(TypeDefinition[] allTypes, TypeDefinition traitForAttributeTypeDefinition)
{
	var result = new Dictionary< TypeDefinition, IList< MethodDefinition>>();

	var implementers = allTypes.Where(type => type.CustomAttributes.Any(attribute => attribute.AttributeType.FullName == traitForAttributeTypeDefinition.FullName));

	foreach(var implementer in implementers)
	{
		var extendedInterafaceType = implementer.CustomAttributes.First(attribute => attribute.AttributeType.FullName == traitForAttributeTypeDefinition.FullName).ConstructorArguments.First().Value as TypeDefinition;
		var extendedInterfaceTypeDefinition = allTypes.First(type => type.FullName == extendedInterafaceType.FullName);
		var extendingMethods = implementer.Methods;
		result[extendedInterfaceTypeDefinition] = extendingMethods.ToList();
	}
	
	return result;
}

We iterate over all classes marked with our attribute. For each class we extract type from attribute indicated which interface we want to boost. Next, we take all methods from class and store them in the dictionary. Her we implicitly assume that there is only one class with extension methods for particular interface, but it is not difficult to change this behaviour.

Let’s see how we extend interfaces:

private void ExtendInterface(TypeDefinition extendedInterface, IList< MethodDefinition> methods)
{
	foreach (var method in methods)
	{
		InjectMethodToInterface(extendedInterface, method);
	}
}

private void InjectMethodToInterface(TypeDefinition extendedInterface, MethodDefinition method)
{
	var newMethod = new MethodDefinition(method.Name, MethodAttributes.Abstract | MethodAttributes.Virtual | MethodAttributes.HideBySig | MethodAttributes.NewSlot, method.ReturnType);
	extendedInterface.Methods.Add(newMethod);
}

For each interface and for each extension method we create new method (with specific attributes) and simply add this method. Methods are empty so we don’t need to bother with body here.

You might ask: how do you know which attributes should you use? If you look into MethodAttributes in Mono.Cecil you will find out that there are lots of them. Well, it is pretty simple: just create interface with a method, compile the code, and then decompile it using ILSpy or any other tool — you will see all attributes which probably are needed.

OK, we know how to extend interface. Let’s see how to inject implementation to concrete classes:

private void ExtendInheritors(TypeDefinition extendedType, IList< MethodDefinition> methods, TypeDefinition[] allTypes)
{
	var inheritors = GetInheritingTypes(extendedType, allTypes);
	foreach (var inheritor in inheritors)
	{
		ExtendInheritor(inheritor, extendedType, methods, allTypes);
	}
}

private void ExtendInheritor(TypeDefinition inheritor, TypeDefinition extendedType, IList< MethodDefinition> methods, TypeDefinition[] allTypes)
{
	foreach (var method in methods)
	{
		InjectMethodToInheritor(inheritor, extendedType, method, allTypes);
	}
}

private void InjectMethodToInheritor(TypeDefinition inheritor, TypeDefinition extendedType, MethodDefinition method, TypeDefinition[] allTypes)
{
	var newMethod = new MethodDefinition(method.Name, MethodAttributes.Public | MethodAttributes.Virtual | MethodAttributes.HideBySig | MethodAttributes.NewSlot | MethodAttributes.NewSlot, method.ReturnType);
	foreach (var instruction in method.Body.Instructions)
	{
		newMethod.Body.Instructions.Add(instruction);
	}
	inheritor.Methods.Add(newMethod);

	FixDescendants(inheritor, allTypes, method);
}

private void FixDescendants(TypeDefinition inheritor, TypeDefinition[] allTypes, MethodDefinition method)
{
	var primaryDescendants = allTypes.Where(type => type.BaseType == inheritor);
	foreach (var descendant in primaryDescendants)
	{
		foreach (var methodToFix in descendant.Methods.Where(descendantMethod => descendantMethod.Name == method.Name))
		{
			methodToFix.Attributes |= MethodAttributes.Virtual;
			methodToFix.Attributes &= ~(MethodAttributes.NewSlot);
		}
		FixDescendants(descendant, allTypes, method);
	}
}

We iterate over all types, over all methods, and create new methods. We do almost the same as when modifying interfaces, but this time we need to add body to methods (and also please note that methods’ attributes are different this time). Since we want to copy extension methods exactly, we simply iterate over all instructions and add them one by one. Finally, we need to fix subtypes — each subtype is allowed to define method with the same signature which does not need to be marked as virtual. We need to take care of that in order to have polymorphic invocation. In order to do that, we simply add attribute Virtual to each method, and remove attribute NewSlot. The former should be pretty obvious. The latter means that this method hides method with the same signature in the base class. If we remove this attribute, polymorphic invocation will work.

Finally, let’s see how to fix the extension methods:

private void FixTraitMethods(TypeDefinition extendedInterface, IList< MethodDefinition> methods)
{
	foreach (var method in methods)
	{
		FixTraitMethod(extendedInterface, method);
	}
}

private void FixTraitMethod(TypeDefinition extendedInterface, MethodDefinition method)
{
	var methodToCall = extendedInterface.Methods.First(interfaceMethod => interfaceMethod.Name == method.Name);
	LogWarning(methodToCall.FullName);
	method.Body.Instructions.Clear();
	method.Body.Instructions.Add(Instruction.Create(OpCodes.Ldarg_0));
	method.Body.Instructions.Add(Instruction.Create(OpCodes.Callvirt, methodToCall));
	method.Body.Instructions.Add(Instruction.Create(OpCodes.Ret));
}

This time we need to write some magic code. Since we need to call method from interface, we first find its definition. Next, we add IL code which performs the invocation. Since we assume that methods are parameterless and void, we only need to push instance on the stack and call the method. That’s all.

Test code and results

Let’s compile our demo program with Fody enabled and see final IL (using DotPeek):

Demo program:

// Type: TraitsDemo.Program
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 884502CE-D2C9-4FE4-BF73-324C59278A5B
// Location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe
// Sequence point data from decompiler

.class private auto ansi beforefieldinit
  TraitsDemo.Program
    extends [mscorlib]System.Object
{

  .method private hidebysig static void
    Main(
      string[] args
    ) cil managed
  {
    .entrypoint
    .maxstack 1
    .locals init (
      [0] class TraitsDemo.IA V_0
    )

    IL_0000: nop         

    // [13 7 - 13 39]
    IL_0001: newobj       instance void TraitsDemo.B::.ctor()
    IL_0006: stloc.0      // V_0
    IL_0007: ldloc.0      // V_0
    IL_0008: call         void TraitsDemo.IA_Implementation::Print(class TraitsDemo.IA)
    IL_000d: nop         
    IL_000e: ret         

  } // end of method Program::Main

  .method public hidebysig specialname rtspecialname instance void
    .ctor() cil managed
  {
    .maxstack 8

    IL_0000: ldarg.0      // this
    IL_0001: call         instance void [mscorlib]System.Object::.ctor()
    IL_0006: nop         
    IL_0007: ret         

  } // end of method Program::.ctor
} // end of class TraitsDemo.Program

Class A:

// Type: TraitsDemo.A
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 884502CE-D2C9-4FE4-BF73-324C59278A5B
// Location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe
// Sequence point data from decompiler

.class public auto ansi beforefieldinit
  TraitsDemo.A
    extends [mscorlib]System.Object
    implements TraitsDemo.IA
{

  .method public hidebysig specialname rtspecialname instance void
    .ctor() cil managed
  {
    .maxstack 8

    IL_0000: ldarg.0      // this
    IL_0001: call         instance void [mscorlib]System.Object::.ctor()
    IL_0006: nop         
    IL_0007: ret         

  } // end of method A::.ctor

  .method public hidebysig virtual newslot instance void
    Print() cil managed
  {
    .maxstack 8

    IL_0000: nop         

    // [15 7 - 15 34]
    IL_0001: ldstr        "I'm IA"
    IL_0006: call         void [mscorlib]System.Console::WriteLine(string)
    IL_000b: nop         
    IL_000c: ret         

  } // end of method A::Print
} // end of class TraitsDemo.A

Interface IA:

// Type: TraitsDemo.IA
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 884502CE-D2C9-4FE4-BF73-324C59278A5B
// Location: C:\Adam\bin\Debug\TraitsDemo.exe
// Sequence point data from decompiler

.class interface public abstract auto ansi
  TraitsDemo.IA
{

  .method compilercontrolled hidebysig virtual newslot abstract instance void
    Print() cil managed
  {
    // Can't find a body
  } // end of method IA::Print
} // end of class TraitsDemo.IA

Class IA_Implementation:

// Type: TraitsDemo.IA_Implementation
// Assembly: TraitsDemo, Version=1.0.0.0, Culture=neutral, PublicKeyToken=null
// MVID: 884502CE-D2C9-4FE4-BF73-324C59278A5B
// Location: C:\Adam\TraitsDemo\bin\Debug\TraitsDemo.exe
// Sequence point data from decompiler

.class public abstract sealed auto ansi beforefieldinit
  TraitsDemo.IA_Implementation
    extends [mscorlib]System.Object
{
  .custom instance void [TraitIntroducer]TraitIntroducer.TraitForAttribute::.ctor(class [mscorlib]System.Type)
    = (
      01 00 0d 54 72 61 69 74 73 44 65 6d 6f 2e 49 41 // ...TraitsDemo.IA
      00 00                                           // ..
    )
    // MetadataClassType(TraitsDemo.IA)
  .custom instance void [mscorlib]System.Runtime.CompilerServices.ExtensionAttribute::.ctor()
    = (01 00 00 00 )

  .method public hidebysig static void
    Print(
      class TraitsDemo.IA 'instance'
    ) cil managed
  {
    .custom instance void [mscorlib]System.Runtime.CompilerServices.ExtensionAttribute::.ctor()
      = (01 00 00 00 )
    .maxstack 8

    // [16 7 - 16 23]
    IL_0000: ldarg.0      // 'instance'
    IL_0001: callvirt     instance void TraitsDemo.IA::Print()
    IL_0006: ret         

  } // end of method IA_Implementation::Print
} // end of class TraitsDemo.IA_Implementation

Looks pretty decent — and it works.

Summary

This is the end of part one. Next time we are going to write some code to analyse inheritance hierarchy in order to be able to override trait methods in interfaces, so we will have Print in IA and IB : IA, and we will copy correct methods to concrete classes. See you next time!

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.

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:

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:

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:

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:

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:

using (BasicStackAllocator allocator = new BasicStackAllocator())
{
   Expression< Func< CustomType>> lambda = () => new CustomType();
   var customObject = allocator.Alloc(lambda);
   /// ...
}

And we want to get the following:

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:

; 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:

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.