From 667b5d743cfb2e71580cd95ccabbfe3509471ddb Mon Sep 17 00:00:00 2001
From: Jeltz <jeltz@auro.re>
Date: Wed, 7 Sep 2022 06:26:02 +0200
Subject: [PATCH] Don't create an anon set for single IP prefixes

---
 compile.ml | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/compile.ml b/compile.ml
index 474be8d..1d21b2e 100644
--- a/compile.ml
+++ b/compile.ml
@@ -1,4 +1,5 @@
 open Ipaddr
+open Utils
 
 module Prefix = struct
   type t = Ipv4 of V4.Prefix.t | Ipv6 of V6.Prefix.t | Not of t
@@ -89,19 +90,25 @@ module Rules = struct
       []
     |> List.map expr
 
+  let wrap_set = function
+    | [] -> None
+    | [ x ] -> Some x
+    | xs -> Some (Expr.Set xs)
+
   let compile_match_addrs getter expr field zones addrs_list =
     let equal = compile_addrs_list getter expr false zones addrs_list in
     let not_equal = compile_addrs_list getter expr true zones addrs_list in
     let stmts =
-      match equal with
-      | [] -> []
-      | _ -> [ Stmt.Match (Match.Eq, Expr.Payload field, Expr.Set equal) ]
+      [
+        Option.map
+          (fun e -> Stmt.Match (Match.Eq, Expr.Payload field, e))
+          (wrap_set equal);
+        Option.map
+          (fun e -> Stmt.Match (Match.NotEq, Expr.Payload field, e))
+          (wrap_set not_equal);
+      ]
     in
-    match not_equal with
-    | [] -> stmts
-    | _ ->
-        Stmt.Match (Match.NotEq, Expr.Payload field, Expr.Set not_equal)
-        :: stmts
+    deoptionalise stmts
 
   let compile_match_ipv4 field =
     compile_match_addrs Prefix.to_ipv4_list Expr.ipv4 (Payload.Ipv4 field)