diff --git a/bitmask/bitmask.go b/bitmask/bitmask.go index 4871909..1cc5ae0 100644 --- a/bitmask/bitmask.go +++ b/bitmask/bitmask.go @@ -24,3 +24,11 @@ func Remove[T Bit](sum, bit T) T { func Has[T Bit](sum, bit T) bool { return (sum & bit) == bit } + +// Toggle will either add or remove bit from sum depending on whether it currently exists in the bitmask. +func Toggle[T Bit](sum, bit T) T { + if Has(sum, bit) { + return Remove(sum, bit) + } + return Add(sum, bit) +} diff --git a/bitmask/bitmask_test.go b/bitmask/bitmask_test.go index cacfda4..9bb8a34 100644 --- a/bitmask/bitmask_test.go +++ b/bitmask/bitmask_test.go @@ -20,4 +20,7 @@ func TestBits(t *testing.T) { assert.Equal(t, a, bitmask.Remove(a|b, b)) assert.True(t, bitmask.Has(a|b, b)) assert.False(t, bitmask.Has(a|b, c)) + + assert.Equal(t, a|b, bitmask.Toggle(a, b)) + assert.Equal(t, b, bitmask.Toggle(a|b, a)) }