From a4720f30e532c6c2815dd8ea626cf371ad2dca0e Mon Sep 17 00:00:00 2001
From: Conner Fromknecht <conner@lightning.engineering>
Date: Fri, 19 Apr 2019 16:26:22 -0700
Subject: [PATCH] txscript: Optimize removeOpcodeRaw

---
 txscript/script.go      | 39 +++++++++++++++++++++++++++++++++++++++
 txscript/script_test.go |  7 +++----
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/txscript/script.go b/txscript/script.go
index 336a2c8b..7bfe44b3 100644
--- a/txscript/script.go
+++ b/txscript/script.go
@@ -269,6 +269,8 @@ func DisasmString(script []byte) (string, error) {
 
 // removeOpcode will remove any opcode matching ``opcode'' from the opcode
 // stream in pkscript
+//
+// DEPRECATED.  Use removeOpcodeRaw instead.
 func removeOpcode(pkscript []parsedOpcode, opcode byte) []parsedOpcode {
 	retScript := make([]parsedOpcode, 0, len(pkscript))
 	for _, pop := range pkscript {
@@ -279,6 +281,43 @@ func removeOpcode(pkscript []parsedOpcode, opcode byte) []parsedOpcode {
 	return retScript
 }
 
+// removeOpcodeRaw will return the script after removing any opcodes that match
+// `opcode`. If the opcode does not appear in script, the original script will
+// be returned unmodified. Otherwise, a new script will be allocated to contain
+// the filtered script. This metehod assumes that the script parses
+// successfully.
+//
+// NOTE: This function is only valid for version 0 scripts.  Since the function
+// does not accept a script version, the results are undefined for other script
+// versions.
+func removeOpcodeRaw(script []byte, opcode byte) []byte {
+	// Avoid work when possible.
+	if len(script) == 0 {
+		return script
+	}
+
+	const scriptVersion = 0
+	var result []byte
+	var prevOffset int32
+
+	tokenizer := MakeScriptTokenizer(scriptVersion, script)
+	for tokenizer.Next() {
+		if tokenizer.Opcode() == opcode {
+			if result == nil {
+				result = make([]byte, 0, len(script))
+				result = append(result, script[:prevOffset]...)
+			}
+		} else if result != nil {
+			result = append(result, script[prevOffset:tokenizer.ByteIndex()]...)
+		}
+		prevOffset = tokenizer.ByteIndex()
+	}
+	if result == nil {
+		return script
+	}
+	return result
+}
+
 // isCanonicalPush returns true if the opcode is either not a push instruction
 // or the data associated with the push instruction uses the smallest
 // instruction to do the job.  False otherwise.
diff --git a/txscript/script_test.go b/txscript/script_test.go
index 9a64865e..02c364ce 100644
--- a/txscript/script_test.go
+++ b/txscript/script_test.go
@@ -3981,13 +3981,12 @@ func TestRemoveOpcodes(t *testing.T) {
 	// tstRemoveOpcode is a convenience function to parse the provided
 	// raw script, remove the passed opcode, then unparse the result back
 	// into a raw script.
+	const scriptVersion = 0
 	tstRemoveOpcode := func(script []byte, opcode byte) ([]byte, error) {
-		pops, err := parseScript(script)
-		if err != nil {
+		if err := checkScriptParses(scriptVersion, script); err != nil {
 			return nil, err
 		}
-		pops = removeOpcode(pops, opcode)
-		return unparseScript(pops)
+		return removeOpcodeRaw(script, opcode), nil
 	}
 
 	for _, test := range tests {