From 2e2bf1d02eadb41567b1bfdc32b4ba5951813682 Mon Sep 17 00:00:00 2001
From: Ryo Kitagawa <kitadrum50@gmail.com>
Date: Tue, 18 Jan 2022 23:05:58 +0900
Subject: [PATCH] Fix completion duplicates

---
 langserver/internal/source/completion.go      | 23 ++++++++++-----
 langserver/internal/source/completion_test.go | 29 +++++++++++++++++++
 2 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/langserver/internal/source/completion.go b/langserver/internal/source/completion.go
index 2cfde7b..bc4385f 100644
--- a/langserver/internal/source/completion.go
+++ b/langserver/internal/source/completion.go
@@ -143,19 +143,15 @@ func (p *Project) listCompletionItemsForTerms(location *ast.Location, target *as
 			result = append(result, list...)
 		}
 
-		for _, r := range module.Rules {
-			result = append(result, createRuleCompletionItem(r))
-		}
+		result = append(result, p.listCompletionItemsModuleRules(module.Rules)...)
 	}
 
 	if p.isLibraryTerm(target) {
 		if _, ok := target.Value.(ast.Ref); ok {
 			importRef := p.findPolicyRef(target)
 			policies := p.cache.FindPolicies(importRef)
-			for _, p := range policies {
-				for _, r := range p.Rules {
-					result = append(result, createRuleCompletionItem(r))
-				}
+			for _, po := range policies {
+				result = append(result, p.listCompletionItemsModuleRules(po.Rules)...)
 			}
 		}
 	}
@@ -238,6 +234,19 @@ func (p *Project) listCompletionItemsInTerm(loc *ast.Location, term *ast.Term) [
 	return result
 }
 
+func (p *Project) listCompletionItemsModuleRules(rules []*ast.Rule) []CompletionItem {
+	result := make([]CompletionItem, 0, len(rules))
+	exists := make(map[string]struct{}, 0)
+	for _, r := range rules {
+		item := createRuleCompletionItem(r)
+		if _, ok := exists[item.Label]; !ok {
+			result = append(result, item)
+			exists[item.Label] = struct{}{}
+		}
+	}
+	return result
+}
+
 func (p *Project) listBuiltinFunction(term *ast.Term) []CompletionItem {
 	if term == nil {
 		return nil
diff --git a/langserver/internal/source/completion_test.go b/langserver/internal/source/completion_test.go
index ad767a0..e550331 100644
--- a/langserver/internal/source/completion_test.go
+++ b/langserver/internal/source/completion_test.go
@@ -89,6 +89,35 @@ authorize = "allow" {
 				},
 			},
 		},
+		"completion other args": {
+			files: map[string]source.File{
+				"src.rego": {
+					RowText: `package src
+
+func() {
+	me
+}
+
+mem_multiple("E") = 1000000000000000000000
+
+mem_multiple("P") = 1000000000000000000`,
+				},
+			},
+			location: &ast.Location{
+				Row: 4,
+				Col: 3,
+				Offset: len("package src\n\nfunc() {\n	me"),
+				Text: []byte("e"),
+				File: "src.rego",
+			},
+			expectItems: []source.CompletionItem{
+				{
+					Label:      "mem_multiple",
+					Kind:       source.FunctionItem,
+					InsertText: `mem_multiple("E")`,
+				},
+			},
+		},
 	}
 
 	for n, tt := range tests {