How to Extend a Type's State by Adding Fields with ConditionalWeakTable

Posted by : on

Category : C#

Introduction

Some time ago, I wrote a post titled: 4 Ways To Use Extension Methods In C# Other Than Extending Existing Types. While extension methods are a great way to add behavior to types that we don’t own, they do not allow us to add new states.

Occasionally, we may find ourselves needing to add fields with data to types that we don’t own. In this post, I will explain how this can be achieved using the ConditionalWeakTable.

The ConditionalWeakTable is a data structure that resembles a dictionary. However, despite the similarity, there are significant differences in both the available methods and how they manage references. I will discuss these differences after providing an example of how it can be used to add state to an existing type.

Example Of Adding State

To begin, let’s assume we have a type that we don’t own, meaning we can’t modify its code. For this example, the Enemy type is provided as-is, and we only have access to use it, not alter its implementation:

public class Enemy
{
   public int AttackPower { get; }
   public Enemy(int attackPower) => AttackPower = attackPower;
   public void Attack() => Console.WriteLine($"Attacking doing {AttackPower} damage.");
}

That’s a simple type, but it will serve well for our example.

This method will perform an attack where the damage is multiplied by a specific value, which we’ll call attackMultiplier. While we could create a method that takes the multiplier as an argument, we want this multiplier to be a property of the type itself. Ideally, we’d like to set this multiplier for each instance of the type, allowing us to check, modify, and use it whenever MultiAttack is called, without needing to pass it as an argument each time.

In essence, we want to extend the type by adding a new field, which will be encapsulated with methods to get, set, and use it.

The first step is to create a class that contains all the new fields we want to add to our type. For example:

public class EnemyData
{
    public int AttackMultiplier;
}

Currently, this class is public. However, we likely don’t want it to be accessible for creating standalone instances, as its sole purpose is to provide additional fields for our Enemy class. To encapsulate it, we can make the class private and nest it within another class that will serve as a container for the extensions to our Enemy type:

public static class EnemyExtensions
{
   private class EnemyData
   {
      public int AttackMultiplier;
   }
}

We can now begin using the ConditionalWeakTable. This is a generic type that takes two parameters: the first serves as the key, and the second serves as the value associated with that key. Unlike a dictionary, the value must be a reference type. In our case, the keys will be instances of the Enemy class, and the values will be instances of the EnemyData class:

public static class EnemyExtensions
{
   private static ConditionalWeakTable<Enemy, EnemyData> enemyDataTable = new();

   private class EnemyData
   {
      public int AttackMultiplier;
   }
}

The ConditionalWeakTable includes several expected methods, such as Add, AddOrUpdate, Remove, and Clear. In addition, it provides two particularly useful methods: GetValue and GetOrCreateValue.

The GetOrCreateValue method takes a single argument, the key. If the key doesn’t already exist in the table, an entry is created with the given key, and the value is instantiated by invoking the parameterless constructor of the reference type.

The GetValue method takes two arguments: the key and a callback function responsible for creating an instance of the value type. This method is useful when the value type does not have, or you do not wish to use, a parameterless constructor

If you only need to retrieve the value for a key without adding a new one, you can use the TryGetValue method. Additionally, instances of ConditionalWeakTable are thread-safe, meaning no explicit locking is required for operations in a multi-threaded environment.

Let’s incorporate the GetOrCreateValue method into our EnemyExtensions type:

public static class EnemyExtensions
{
   private static ConditionalWeakTable<Enemy, EnemyData> enemyDataTable = new();

   private static EnemyData Data(Enemy enemy) => enemyDataTable.GetOrCreateValue(enemy);

   private class EnemyData
   {
      public int AttackMultiplier;
   }
}

We’ve completed most of the work. Now, each Enemy instance will automatically be paired with an EnemyData instance whenever the Data method is invoked.

While we could make this method public, it’s better to encapsulate it and provide getters and setters for the relevant fields. Unfortunately, C# does not currently support extension properties. Although such a feature has been under discussion for some time, we don’t yet have a mechanism for it. Until then, we must rely on the traditional approach: creating a method to retrieve the value and another to set it.

public static class EnemyExtensions
{
   private static ConditionalWeakTable<Enemy, EnemyData> enemyDataTable = new();

   public static void SetAttackMultiplier(this Enemy enemy, int multiplier) => Data(enemy).AttackMultiplier = multiplier;
   public static int GetAttackMultiplier(this Enemy enemy) => Data(enemy).AttackMultiplier;

   private static EnemyData Data(Enemy enemy) => enemyDataTable.GetOrCreateValue(enemy);

   private class EnemyData
   {
      public int AttackMultiplier;
   }
}

We’re almost finished. We’ve successfully added an extension field to our class, and the final step is to create an extension method that utilizes this field. This process is straightforward and follows the standard approach for creating extension methods. Here’s the complete code:

public static class EnemyExtensions
{
   private static ConditionalWeakTable<Enemy, EnemyData> enemyDataTable = new();

   public static void SetAttackMultiplier(this Enemy enemy, int multiplier) => Data(enemy).AttackMultiplier = multiplier;
   public static int GetAttackMultiplier(this Enemy enemy) => Data(enemy).AttackMultiplier;

   public static void MultiAttack(this Enemy enemy)
   {
      int attackMultiplier = enemy.GetAttackMultiplier();

      Console.WriteLine(attackMultiplier > 0
         ? $"Attacking {attackMultiplier} times doing in total {attackMultiplier * enemy.AttackPower} damage."
         : $"Cannot Attack multiple times. Performing normal attack doing {enemy.AttackPower} damage.");
   }
   
   private static EnemyData Data(Enemy enemy) => enemyDataTable.GetOrCreateValue(enemy);

   private class EnemyData
   {
      public int AttackMultiplier;
   }
}

Now we can write code like this:

Enemy enemy = new(10);
Enemy enemy2 = new(100);

enemy.Attack();
enemy.MultiAttack();
Console.WriteLine($"The enemy's attack multiplier before setting it is {enemy.GetAttackMultiplier()}");
enemy.SetAttackMultiplier(5);
Console.WriteLine($"The enemy's attack multiplier after setting it is {enemy.GetAttackMultiplier()}");
enemy.Attack();
enemy.MultiAttack();
Console.WriteLine();
enemy2.MultiAttack();
enemy2.SetAttackMultiplier(2);
enemy.MultiAttack();
enemy2.MultiAttack();

Our Enemy class has behaves for its users like having a new field, that we can add or change its value and also has new functionality that utilizes this field.

Differences Between Using a Dictionary and a ConditionalWeakTable

The previous example could have been implemented using a Dictionary instead of a ConditionalWeakTable, but that approach would be incorrect. While there are several differences between these two types, most are minor inconveniences. However, the key distinction is that the ConditionalWeakTable holds a WeakReference to its keys, which significantly impacts memory usage and how the garbage collector handles Enemy instances.

Unlike a Dictionary, a ConditionalWeakTable does not implement the IDictionary interface and lacks methods like GetEnumerator or Contains. Most importantly, if a key instance is no longer referenced anywhere else in the code, it will be garbage collected. The ConditionalWeakTable will not keep it alive merely because it exists as a key, making it more memory-efficient in scenarios where keys may become unreferenced.

For example:

object? dictKey = new();
object? weakKey = new();
object dictValue = new();
object weakValue = new();

Dictionary<object, object> dict = new();
ConditionalWeakTable<object, object> weakTable = new();

dict.Add(dictKey, dictValue);
weakTable.Add(weakKey,weakValue);

Console.WriteLine(dict.Count);
Console.WriteLine(weakTable.Count());

dictKey = null;
weakKey = null;

await Task.Delay(1000);
GC.Collect();

Console.WriteLine(dict.Count);
Console.WriteLine(weakTable.Count());

The result of running this code will be:

1
1
1
0

The instance that has been created and assigned to weakKey won’t be kept alive.

Conclusion

This is how we can currently extend a type we don’t own with additional fields. Until (or if?) we get a unified language feature that allows us to extend types with more than just methods, something that has been discussed in the past, we can use the ConditionalWeakTable along with a class that holds the desired fields.

As always, thank you for reading, and if you have any questions or comments you can use the comments section, or contact me directly via the contact form or by email. Also, if you don’t want to miss any of the new blog posts, you can subscribe to my newsletter or the RSS feed.


Follow me: