From 55eb898186a0dcaa80eae5adced7f0bd16abea16 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Fri, 25 Nov 2022 13:37:50 +0100
Subject: [PATCH] Add support for specifying multiple prefixes

This change allows the file Locator to be instantiated with multiple
search prefixes.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 internal/lookup/file.go      | 46 +++++++++++++++++++++++++++++++-----
 internal/lookup/file_test.go | 41 +++++++++++++++++++++++++++-----
 2 files changed, 75 insertions(+), 12 deletions(-)

diff --git a/internal/lookup/file.go b/internal/lookup/file.go
index ab52d03d..724075a2 100644
--- a/internal/lookup/file.go
+++ b/internal/lookup/file.go
@@ -32,22 +32,56 @@ type file struct {
 	filter   func(string) error
 }
 
-// NewFileLocator creates a Locator that can be used to find files at the specified root. A logger
-// can also be specified.
-func NewFileLocator(logger *log.Logger, root string) Locator {
-	l := newFileLocator(logger, root)
+// NewFileLocator creates a Locator that can be used to find files at the specified root.
+// An optional list of prefixes can aslo be specified with each of these being searched in order.
+// The specified root is prefixed to each of the prefixes to determine the final search path.
+func NewFileLocator(logger *log.Logger, root string, prefixes ...string) Locator {
+	l := newFileLocator(logger, root, prefixes...)
 
 	return &l
 }
 
-func newFileLocator(logger *log.Logger, root string) file {
+func newFileLocator(logger *log.Logger, root string, prefixes ...string) file {
+
 	return file{
 		logger:   logger,
-		prefixes: []string{root},
+		prefixes: getSearchPrefixes(root, prefixes...),
 		filter:   assertFile,
 	}
 }
 
+// getSearchPrefixes generates a list of unique paths to be searched by a file locator.
+//
+// For each of the unique prefixes <p> specified the path <root><p> is searched, where <root> is the
+// specified root. If no prefixes are specified, <root> is returned as the only search prefix.
+//
+// Note that an empty root is equivalent to searching relative to the current working directory, and
+// if the root filesystem should be searched instead, root should be specified as "/" explicitly.
+//
+// Also, a prefix of "" forces the root to be included in returned set of paths. This means that if
+// the root in addition to another prefix must be searched the function should be called with:
+//
+//	getSearchPrefixes("/root", "", "another/path")
+//
+// and will result in the search paths []{"/root", "/root/another/path"} being returned.
+func getSearchPrefixes(root string, prefixes ...string) []string {
+	seen := make(map[string]bool)
+	var uniquePrefixes []string
+	for _, p := range prefixes {
+		if seen[p] {
+			continue
+		}
+		seen[p] = true
+		uniquePrefixes = append(uniquePrefixes, filepath.Join(root, p))
+	}
+
+	if len(uniquePrefixes) == 0 {
+		uniquePrefixes = append(uniquePrefixes, root)
+	}
+
+	return uniquePrefixes
+}
+
 var _ Locator = (*file)(nil)
 
 // Locate attempts to find files with names matching the specified pattern.
diff --git a/internal/lookup/file_test.go b/internal/lookup/file_test.go
index 78285554..b7006bf5 100644
--- a/internal/lookup/file_test.go
+++ b/internal/lookup/file_test.go
@@ -20,18 +20,17 @@ import (
 	"fmt"
 	"testing"
 
-	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
 )
 
-func TestFileLocator(t *testing.T) {
-	logger, _ := testlog.NewNullLogger()
-
+func TestGetSearchPrefixes(t *testing.T) {
 	testCases := []struct {
 		root             string
+		prefixes         []string
 		expectedPrefixes []string
 	}{
 		{
+			root:             "",
 			expectedPrefixes: []string{""},
 		},
 		{
@@ -42,12 +41,42 @@ func TestFileLocator(t *testing.T) {
 			root:             "/some/root",
 			expectedPrefixes: []string{"/some/root"},
 		},
+		{
+			root:             "",
+			prefixes:         []string{"foo", "bar"},
+			expectedPrefixes: []string{"foo", "bar"},
+		},
+		{
+			root:             "/",
+			prefixes:         []string{"foo", "bar"},
+			expectedPrefixes: []string{"/foo", "/bar"},
+		},
+		{
+			root:             "/",
+			prefixes:         []string{"/foo", "/bar"},
+			expectedPrefixes: []string{"/foo", "/bar"},
+		},
+		{
+			root:             "/some/root",
+			prefixes:         []string{"foo", "bar"},
+			expectedPrefixes: []string{"/some/root/foo", "/some/root/bar"},
+		},
+		{
+			root:             "",
+			prefixes:         []string{"foo", "bar", "bar", "foo"},
+			expectedPrefixes: []string{"foo", "bar"},
+		},
+		{
+			root:             "/some/root",
+			prefixes:         []string{"foo", "bar", "foo", "bar"},
+			expectedPrefixes: []string{"/some/root/foo", "/some/root/bar"},
+		},
 	}
 
 	for i, tc := range testCases {
 		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
-			f := newFileLocator(logger, tc.root)
-			require.EqualValues(t, tc.expectedPrefixes, f.prefixes)
+			prefixes := getSearchPrefixes(tc.root, tc.prefixes...)
+			require.EqualValues(t, tc.expectedPrefixes, prefixes)
 		})
 	}
 }