Merge pull request #1954 from guggero/script-builder-alloc

txscript: allow script builder capacity to be specified

txscript/scriptbuilder.go

@@ -13,11 +13,41 @@ const (
 	// defaultScriptAlloc is the default size used for the backing array
 	// for a script being built by the ScriptBuilder.  The array will
 	// dynamically grow as needed, but this figure is intended to provide
-	// enough space for vast majority of scripts without needing to grow the
-	// backing array multiple times.
+	// enough space for the vast majority of scripts without needing to grow
+	// the backing array multiple times. Can be overwritten with the
+	// WithScriptAllocSize functional option where expected script sizes are
+	// known.
 	defaultScriptAlloc = 500
 )
 
+// scriptBuilderConfig is a configuration struct that can be used to modify the
+// initialization of a ScriptBuilder.
+type scriptBuilderConfig struct {
+	// allocSize specifies the initial size of the backing array for the
+	// script builder.
+	allocSize int
+}
+
+// defaultScriptBuilderConfig returns a new scriptBuilderConfig with the
+// default values set.
+func defaultScriptBuilderConfig() *scriptBuilderConfig {
+	return &scriptBuilderConfig{
+		allocSize: defaultScriptAlloc,
+	}
+}
+
+// ScriptBuilderOpt is a functional option type which is used to modify the
+// initialization of a ScriptBuilder.
+type ScriptBuilderOpt func(*scriptBuilderConfig)
+
+// WithScriptAllocSize specifies the initial size of the backing array for the
+// script builder.
+func WithScriptAllocSize(size int) ScriptBuilderOpt {
+	return func(cfg *scriptBuilderConfig) {
+		cfg.allocSize = size
+	}
+}
+
 // ErrScriptNotCanonical identifies a non-canonical script.  The caller can use
 // a type assertion to detect this error type.
 type ErrScriptNotCanonical string
@@ -37,16 +67,17 @@ func (e ErrScriptNotCanonical) Error() string {
 // For example, the following would build a 2-of-3 multisig script for usage in
 // a pay-to-script-hash (although in this situation MultiSigScript() would be a
 // better choice to generate the script):
-// 	builder := txscript.NewScriptBuilder()
-// 	builder.AddOp(txscript.OP_2).AddData(pubKey1).AddData(pubKey2)
-// 	builder.AddData(pubKey3).AddOp(txscript.OP_3)
-// 	builder.AddOp(txscript.OP_CHECKMULTISIG)
-// 	script, err := builder.Script()
-// 	if err != nil {
-// 		// Handle the error.
-// 		return
-// 	}
-// 	fmt.Printf("Final multi-sig script: %x\n", script)
+//
+//	builder := txscript.NewScriptBuilder()
+//	builder.AddOp(txscript.OP_2).AddData(pubKey1).AddData(pubKey2)
+//	builder.AddData(pubKey3).AddOp(txscript.OP_3)
+//	builder.AddOp(txscript.OP_CHECKMULTISIG)
+//	script, err := builder.Script()
+//	if err != nil {
+//		// Handle the error.
+//		return
+//	}
+//	fmt.Printf("Final multi-sig script: %x\n", script)
 type ScriptBuilder struct {
 	script []byte
 	err    error
@@ -267,8 +298,13 @@ func (b *ScriptBuilder) Script() ([]byte, error) {
 
 // NewScriptBuilder returns a new instance of a script builder.  See
 // ScriptBuilder for details.
-func NewScriptBuilder() *ScriptBuilder {
+func NewScriptBuilder(opts ...ScriptBuilderOpt) *ScriptBuilder {
+	cfg := defaultScriptBuilderConfig()
+	for _, opt := range opts {
+		opt(cfg)
+	}
+
 	return &ScriptBuilder{
-		script: make([]byte, 0, defaultScriptAlloc),
+		script: make([]byte, 0, cfg.allocSize),
 	}
 }

txscript/scriptbuilder_test.go

@@ -7,8 +7,38 @@ package txscript
 import (
 	"bytes"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
+// TestScriptBuilderAlloc tests that the pre-allocation for a script via the
+// NewScriptBuilder function works as expected.
+func TestScriptBuilderAlloc(t *testing.T) {
+	// Using the default value, we should get a script with a capacity of
+	// 500 bytes, which is quite large for most scripts.
+	defaultBuilder := NewScriptBuilder()
+	require.EqualValues(t, defaultScriptAlloc, cap(defaultBuilder.script))
+
+	const allocSize = 23
+	builder := NewScriptBuilder(WithScriptAllocSize(allocSize))
+
+	// The initial capacity of the script should be set to the explicit
+	// value.
+	require.EqualValues(t, allocSize, cap(builder.script))
+
+	builder.AddOp(OP_HASH160)
+	builder.AddData(make([]byte, 20))
+	builder.AddOp(OP_EQUAL)
+	script, err := builder.Script()
+	require.NoError(t, err)
+
+	require.Len(t, script, allocSize)
+
+	// The capacity shouldn't have changed, as the script should've fit just
+	// fine.
+	require.EqualValues(t, allocSize, cap(builder.script))
+}
+
 // TestScriptBuilderAddOp tests that pushing opcodes to a script via the
 // ScriptBuilder API works as expected.
 func TestScriptBuilderAddOp(t *testing.T) {