This is the third part of 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