From ffcf6ba6343ef55855a503066488b4c93533d272 Mon Sep 17 00:00:00 2001
From: HikariKnight <2557889+HikariKnight@users.noreply.github.com>
Date: Wed, 26 Apr 2023 18:00:31 +0200
Subject: [PATCH] change -r to be look for related devices in the same group
 and -rr to look for group and vendorID

---
 pkg/iommu/iommu.go   | 26 +++++++++++++-------------
 pkg/params/params.go |  4 ++--
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/pkg/iommu/iommu.go b/pkg/iommu/iommu.go
index fc23397..4426a47 100644
--- a/pkg/iommu/iommu.go
+++ b/pkg/iommu/iommu.go
@@ -139,7 +139,7 @@ func MatchSubclass(searchval string, pArg *params.Params) []string {
 	alldevs := NewIOMMU()
 
 	// Make a regex to find devices that will need special exceptions for relative search
-	specialRelativeExceptions := regexp.MustCompile(`^(SATA|USB) controller`)
+	//specialRelativeExceptions := regexp.MustCompile(`^(SATA|USB) controller`)
 
 	// Iterate through the groups
 	for id := 0; id < len(alldevs.Groups); id++ {
@@ -153,15 +153,15 @@ func MatchSubclass(searchval string, pArg *params.Params) []string {
 					devs = append(devs, line)
 
 					// If we want to search for related devices
-					if pArg.FlagCounter["related"] > 0 && !specialRelativeExceptions.MatchString(searchval) {
-						// Find relatives and add them to the list
-						related_list := findRelatedDevices(device.Vendor.ID, pArg.FlagCounter["related"], pArg)
-						devs = append(devs, related_list...)
-
-					} else if pArg.FlagCounter["related"] > 0 && specialRelativeExceptions.MatchString(searchval) {
+					if pArg.FlagCounter["related"] == 1 {
 						// Prevent an infinite loop by passing 0 instead of related
 						other := GetDevicesFromGroups([]int{id}, 0, pArg)
 						devs = append(devs, other...)
+
+					} else if pArg.FlagCounter["related"] == 2 {
+						// Find relatives and add them to the list
+						related_list := findRelatedDevices(device.Vendor.ID, pArg.FlagCounter["related"], pArg)
+						devs = append(devs, related_list...)
 					}
 
 				} else if pArg.Flag["rom"] && pArg.Flag["gpu"] {
@@ -200,15 +200,15 @@ func MatchSubclass(searchval string, pArg *params.Params) []string {
 							}
 
 							// If we want to search for related devices
-							if pArg.FlagCounter["related"] > 0 && !specialRelativeExceptions.MatchString(searchval) {
-								// Find relatives and add them to the list
-								related_list := findRelatedDevices(device.Vendor.ID, pArg.FlagCounter["related"], pArg)
-								devs = append(devs, related_list...)
-
-							} else if pArg.FlagCounter["related"] > 0 && specialRelativeExceptions.MatchString(searchval) {
+							if pArg.FlagCounter["related"] == 1 {
 								// Prevent an infinite loop by passing 0 instead of related
 								other := GetDevicesFromGroups([]int{id}, 0, pArg)
 								devs = append(devs, other...)
+
+							} else if pArg.FlagCounter["related"] == 2 {
+								// Find relatives and add them to the list
+								related_list := findRelatedDevices(device.Vendor.ID, pArg.FlagCounter["related"], pArg)
+								devs = append(devs, related_list...)
 							}
 						}
 					}
diff --git a/pkg/params/params.go b/pkg/params/params.go
index 88e6bbd..064ec8e 100644
--- a/pkg/params/params.go
+++ b/pkg/params/params.go
@@ -75,12 +75,12 @@ func NewParams() *Params {
 
 	related := parser.FlagCounter("r", "related", &argparse.Options{
 		Required: false,
-		Help:     "Attempt to list related devices that share Vendor ID or\n\t\t IOMMU Groups (used with -g -u -i and -n), pass -rr if you want to search using both when used with -g -i or -n\n\t\t Note: -rr can be inaccurate or too broad when many devices share Vendor ID",
+		Help:     "Attempt to list related devices that share IOMMU Groups or\n\t\t Vendor ID (used with -g -u -i -s and -n), pass -rr if you want to search using both when used with -g -i -s or -n\n\t\t Note: -rr can be inaccurate or too broad when many devices share Vendor ID",
 	})
 
 	ignore := parser.StringList("R", "ignore", &argparse.Options{
 		Required: false,
-		Help:     "Ignores passed VendorID (Left part of : in [VendorID:DeviceID]) outside of the selected IOMMU group when doing a --related search, you can use this to ignore unreliable Vendor IDs when doing related searches. (works with -g -i and -n)",
+		Help:     "Ignores passed VendorID (Left part of : in [VendorID:DeviceID]) outside of the selected IOMMU group when doing a --related search, you can use this to ignore unreliable Vendor IDs when doing related searches. (works with -g -i -u -s and -n)",
 	})
 
 	kernelmodules := parser.Flag("k", "kernel", &argparse.Options{