From c7ae5be4063bdc71df9a09103a456aaebeac6096 Mon Sep 17 00:00:00 2001
From: Remy Suen <remy.suen@gmail.com>
Date: Mon, 17 Jul 2017 05:08:21 +0900
Subject: [PATCH] Fix #44 Use prefix to suggest relevant build stages only

Signed-off-by: Remy Suen <remy.suen@gmail.com>
---
 src/dockerAssist.ts       | 39 +++++++++++++++++++++++----------------
 test/dockerAssist.test.ts |  6 ++++++
 2 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/src/dockerAssist.ts b/src/dockerAssist.ts
index b787f2c..8a37b61 100644
--- a/src/dockerAssist.ts
+++ b/src/dockerAssist.ts
@@ -9,6 +9,7 @@ import {
 	CompletionItem, CompletionItemKind, InsertTextFormat
 } from 'vscode-languageserver';
 import { Util, KEYWORDS, DIRECTIVE_ESCAPE } from './docker';
+import { Dockerfile } from './parser/dockerfile';
 import { DockerfileParser } from './parser/dockerfileParser';
 import { Copy } from './parser/instructions/copy';
 
@@ -75,22 +76,7 @@ export class DockerAssist {
 			} else if (Util.isInsideRange(position, instruction.getRange())) {
 				switch (instruction.getKeyword()) {
 					case "COPY":
-						let copy = instruction as Copy;
-						let copyArgs = instruction.getArguments();
-						if (copyArgs.length !== 0 && copyArgs[0].getValue().indexOf("--from=") === 0 && copy.getFromValueRange().start.character === position.character) {
-							let items: CompletionItem[] = [];
-							for (let from of dockerfile.getFROMs()) {
-								let stage = from.getBuildStage();
-								if (stage) {
-									items.push(this.createSourceImageCompletionItem(stage, "", offset));
-								}
-							}
-							items.sort((item: CompletionItem, item2: CompletionItem) => {
-								return item.label.localeCompare(item2.label);
-							});
-							return items;
-						}
-						return [];
+						return this.createBuildStageProposals(dockerfile, instruction as Copy, position, offset);
 					case "ONBUILD":
 						let onbuildArgs = instruction.getArguments();
 						if (onbuildArgs.length === 0 || Util.isInsideRange(position, onbuildArgs[0].getRange())) {
@@ -160,6 +146,27 @@ export class DockerAssist {
 		return proposals;
 	}
 
+	private createBuildStageProposals(dockerfile: Dockerfile, copy: Copy, position: Position, offset: number) {
+		let range = copy.getFromValueRange();
+		// is the user in the --from= area
+		if (range && Util.isInsideRange(position, copy.getFromValueRange())) {
+			// get the prefix
+			let prefix = this.document.getText().substring(this.document.offsetAt(range.start), offset);
+			let items: CompletionItem[] = [];
+			for (let from of dockerfile.getFROMs()) {
+				let stage = from.getBuildStage();
+				if (stage && stage.indexOf(prefix) === 0) {
+					items.push(this.createSourceImageCompletionItem(stage, prefix, offset));
+				}
+			}
+			items.sort((item: CompletionItem, item2: CompletionItem) => {
+				return item.label.localeCompare(item2.label);
+			});
+			return items;
+		}
+		return [];
+	}
+
 	/**
 	 * Walks back in the text buffer to calculate the true prefix of the
 	 * current text caret offset. This function will handle the
diff --git a/test/dockerAssist.test.ts b/test/dockerAssist.test.ts
index 0c7eb27..7e40f76 100644
--- a/test/dockerAssist.test.ts
+++ b/test/dockerAssist.test.ts
@@ -1236,6 +1236,12 @@ describe('Docker Content Assist Tests', function() {
 				assertSourceImage(proposals[0], "dev", 2, 12, 2, 12);
 				assertSourceImage(proposals[1], "setup", 2, 12, 2, 12);
 			});
+
+			it("source image prefix", function() {
+				var proposals = computePosition("FROM busybox AS setup\nFROM busybox AS dev\nCOPY --from=s", 2, 13);
+				assert.equal(proposals.length, 1);
+				assertSourceImage(proposals[0], "setup", 2, 12, 2, 13);
+			});
 		});
 	});