From b11ccf67be8288f7eb283ca963a7479e2cd1af98 Mon Sep 17 00:00:00 2001
From: GabrielKS <23368820+GabrielKS@users.noreply.github.com>
Date: Wed, 17 Jul 2024 15:07:38 -0600
Subject: [PATCH] `ComponentSelector`: Implement `in(component, selector)`

---
 src/component_selector.jl       | 40 +++++++++++++++++++++++++++++++++
 test/test_component_selector.jl | 24 ++++++++++++++++++++
 2 files changed, 64 insertions(+)

diff --git a/src/component_selector.jl b/src/component_selector.jl
index 488f6a4ba..a7fb18587 100644
--- a/src/component_selector.jl
+++ b/src/component_selector.jl
@@ -4,6 +4,10 @@
 Concrete subtypes MUST implement:
  - `get_components`
  - `get_name`, or use the default by implementing `default_name` and having a `name` field
+ - `Base.in`, with the invariant that for all systems `s` and all `ComponentSelector`s `e`,
+   `c in e` iff `c in collect(get_components(e, s))` for all components `c` in `s`. Ideally,
+   `in` should be evaluable without a reference to a system, but a three-parameter form
+   in(c, e, s) may be used if absolutely necessary.
 
 Concrete subtypes MAY implement:
  - The factory method `select_components` (make sure your signature does not conflict with
@@ -117,6 +121,15 @@ function get_components(e::SingleComponentSelector, sys::Components; filterby =
     return (com === nothing) ? [] : [com]
 end
 
+function Base.in(
+    c::InfrastructureSystemsComponent,
+    e::SingleComponentSelector;
+    filterby = nothing,
+)
+    (!isnothing(filterby) && !filterby(c)) && return false
+    return (c isa e.component_subtype) && (get_name(c) == e.component_name)
+end
+
 # ListComponentSelector
 "ComponentSelectorSet represented by a list of other ComponentSelectors."
 struct ListComponentSelector <: ComponentSelectorSet
@@ -148,6 +161,15 @@ function get_components(e::ListComponentSelector, sys::Components; filterby = no
     return Iterators.flatten(sub_components)
 end
 
+function Base.in(
+    c::InfrastructureSystemsComponent,
+    e::ListComponentSelector;
+    filterby = nothing,
+)
+    (!isnothing(filterby) && !filterby(c)) && return false
+    return any(in.(Ref(c), e.content))
+end
+
 # SubtypeComponentSelector
 "ComponentSelectorSet represented by a subtype of Component."
 struct SubtypeComponentSelector <: ComponentSelectorSet
@@ -181,6 +203,15 @@ function get_components(e::SubtypeComponentSelector, sys::Components; filterby =
     return Iterators.filter(filterby, components)
 end
 
+function Base.in(
+    c::InfrastructureSystemsComponent,
+    e::SubtypeComponentSelector;
+    filterby = nothing,
+)
+    (!isnothing(filterby) && !filterby(c)) && return false
+    return c isa e.component_subtype
+end
+
 # FilterComponentSelector
 "ComponentSelectorSet represented by a filter function and a subtype of Component."
 struct FilterComponentSelector <: ComponentSelectorSet
@@ -225,6 +256,15 @@ function get_components(e::FilterComponentSelector, sys::Components; filterby =
     return Iterators.filter(filterby, components)
 end
 
+function Base.in(
+    c::InfrastructureSystemsComponent,
+    e::FilterComponentSelector;
+    filterby = nothing,
+)
+    (!isnothing(filterby) && !filterby(c)) && return false
+    return (c isa e.component_subtype) && e.filter_fn(c)
+end
+
 # Naming
 default_name(e::FilterComponentSelector) =
     string(e.filter_fn) * NAME_DELIMETER * subtype_to_string(e.component_subtype)
diff --git a/test/test_component_selector.jl b/test/test_component_selector.jl
index a0af93cca..e9602d36d 100644
--- a/test/test_component_selector.jl
+++ b/test/test_component_selector.jl
@@ -71,6 +71,10 @@ end
         @test length(
             collect(IS.get_components(test_gen_ent, test_sys; filterby = x -> false)),
         ) == 0
+
+        @test IS.TestComponent("Component1", -1) in test_gen_ent
+        @test !(IS.TestComponent("Component2", -1) in test_gen_ent)
+        @test !in(IS.TestComponent("Component1", -1), test_gen_ent; filterby = x -> false)
     end
 end
 
@@ -124,6 +128,14 @@ end
         @test length(
             collect(IS.get_subselectors(test_list_ent, test_sys; filterby = x -> false)),
         ) == 2
+
+        @test IS.AdditionalTestComponent("Component3", -1) in test_list_ent
+        @test !(IS.AdditionalTestComponent("Component4", -1) in test_list_ent)
+        @test !in(
+            IS.AdditionalTestComponent("Component3", -1),
+            test_list_ent;
+            filterby = x -> false,
+        )
     end
 end
 
@@ -173,6 +185,10 @@ end
         @test length(
             collect(IS.get_subselectors(test_sub_ent, test_sys; filterby = x -> false)),
         ) == 0
+
+        @test IS.TestComponent("Component1", -1) in test_sub_ent
+        @test !(IS.AdditionalTestComponent("Component1", -1) in test_sub_ent)
+        @test !in(IS.TestComponent("Component1", -1), test_sub_ent; filterby = x -> false)
     end
 end
 
@@ -255,5 +271,13 @@ end
         @test length(
             collect(IS.get_subselectors(test_filter_ent, test_sys; filterby = x -> false)),
         ) == 0
+
+        @test IS.TestComponent("Component1", 20) in test_filter_ent
+        @test !(IS.AdditionalTestComponent("Component1", 0) in test_filter_ent)
+        @test !in(
+            IS.TestComponent("Component1", 20),
+            test_filter_ent;
+            filterby = x -> false,
+        )
     end
 end