From 5c29896686a83b62cb73bf505973fd61227dadc1 Mon Sep 17 00:00:00 2001
From: Tony Holdstock-Brown <tonyhb@gmail.com>
Date: Sun, 7 Jan 2024 21:23:23 -0800
Subject: [PATCH] Add null removal

---
 expr.go      | 28 ++++++++++++++++++++++------
 expr_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/expr.go b/expr.go
index 1f1b938..fa88979 100644
--- a/expr.go
+++ b/expr.go
@@ -158,7 +158,7 @@ func (a *aggregator) Evaluate(ctx context.Context, data map[string]any) ([]Evalu
 		err = errors.Join(err, merr)
 	}
 
-	// TODO: Each match here is a potential success.  When other trees and operators which are walkable
+	// Each match here is a potential success.  When other trees and operators which are walkable
 	// are added (eg. >= operators on strings), ensure that we find the correct number of matches
 	// for each group ID and then skip evaluating expressions if the number of matches is <= the group
 	// ID's length.
@@ -197,8 +197,12 @@ func (a *aggregator) AggregateMatch(ctx context.Context, data map[string]any) ([
 	a.lock.RLock()
 	defer a.lock.RUnlock()
 
-	// Store the number of times each GroupID has found a match.  We need at least
-	// as many matches as stored in the group ID to consider the match.
+	// Each match here is a potential success.  Ensure that we find the correct number of matches
+	// for each group ID and then skip evaluating expressions if the number of matches is <= the group
+	// ID's length.  For example, (A && B && C) is a single group ID and must have a count >= 3,
+	// else we know a required comparason did not match.
+	//
+	// Note that having a count >= the group ID value does not guarantee that the expression is valid.
 	counts := map[groupID]int{}
 	// Store all expression parts per group ID for returning.
 	found := map[groupID][]ExpressionPart{}
@@ -508,7 +512,7 @@ func (a *aggregator) removeNode(ctx context.Context, n *Node, parsed *ParsedExpr
 	case TreeTypeART:
 		tree, ok := a.artIdents[n.Predicate.Ident]
 		if !ok {
-			tree = newArtTree()
+			return ErrExpressionPartNotFound
 		}
 		err := tree.Remove(ctx, ExpressionPart{
 			GroupID:   n.GroupID,
@@ -521,8 +525,20 @@ func (a *aggregator) removeNode(ctx context.Context, n *Node, parsed *ParsedExpr
 		a.artIdents[n.Predicate.Ident] = tree
 		return nil
 	case TreeTypeNullMatch:
-		// TODO: Implement null matching.
-		return errTreeUnimplemented
+		tree, ok := a.nullLookups[n.Predicate.Ident]
+		if !ok {
+			return ErrExpressionPartNotFound
+		}
+		err := tree.Remove(ctx, ExpressionPart{
+			GroupID:   n.GroupID,
+			Predicate: *n.Predicate,
+			Parsed:    parsed,
+		})
+		if err != nil {
+			return err
+		}
+		a.nullLookups[n.Predicate.Ident] = tree
+		return nil
 	}
 	return errTreeUnimplemented
 }
diff --git a/expr_test.go b/expr_test.go
index 903c255..4dcf7ab 100644
--- a/expr_test.go
+++ b/expr_test.go
@@ -530,6 +530,48 @@ func TestEvaluate_Null(t *testing.T) {
 		require.EqualValues(t, 1, count)
 		require.EqualValues(t, isNull, eval[0])
 	})
+
+	t.Run("It removes null checks", func(t *testing.T) {
+		err := e.Remove(ctx, notNull)
+		require.NoError(t, err)
+
+		require.Equal(t, 1, e.Len())
+		require.Equal(t, 0, e.ConstantLen())
+		require.Equal(t, 1, e.AggregateableLen())
+
+		// We should still match on `isNull`
+		t.Run("Is null checks succeed", func(t *testing.T) {
+			// Matching this expr should work, as "ts" is nil
+			eval, count, err := e.Evaluate(ctx, map[string]any{
+				"event": map[string]any{
+					"ts": nil,
+				},
+			})
+			require.NoError(t, err)
+			require.EqualValues(t, 1, len(eval))
+			require.EqualValues(t, 1, count)
+			require.EqualValues(t, isNull, eval[0])
+		})
+
+		err = e.Remove(ctx, isNull)
+		require.NoError(t, err)
+		require.Equal(t, 0, e.Len())
+		require.Equal(t, 0, e.ConstantLen())
+		require.Equal(t, 0, e.AggregateableLen())
+
+		// We should no longer match on `isNull`
+		t.Run("Is null checks succeed", func(t *testing.T) {
+			// Matching this expr should work, as "ts" is nil
+			eval, count, err := e.Evaluate(ctx, map[string]any{
+				"event": map[string]any{
+					"ts": nil,
+				},
+			})
+			require.NoError(t, err)
+			require.EqualValues(t, 0, len(eval))
+			require.EqualValues(t, 0, count)
+		})
+	})
 }
 
 // tex represents a test Evaluable expression