diff --git a/txscript/scriptbuilder.go b/txscript/scriptbuilder.go index 7984dd9661..fa2bc073f0 100644 --- a/txscript/scriptbuilder.go +++ b/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), } } diff --git a/txscript/scriptbuilder_test.go b/txscript/scriptbuilder_test.go index 89f2b861ab..baf9526a26 100644 --- a/txscript/scriptbuilder_test.go +++ b/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) {