github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/protocol/vm/vmutil/builder.go (about)

     1  package vmutil
     2  
     3  import (
     4  	"encoding/binary"
     5  
     6  	"github.com/bytom/bytom/errors"
     7  	"github.com/bytom/bytom/protocol/vm"
     8  )
     9  
    10  type Builder struct {
    11  	program     []byte
    12  	jumpCounter int
    13  
    14  	// Maps a jump target number to its absolute address.
    15  	jumpAddr map[int]uint32
    16  
    17  	// Maps a jump target number to the list of places where its
    18  	// absolute address must be filled in once known.
    19  	jumpPlaceholders map[int][]int
    20  }
    21  
    22  func NewBuilder() *Builder {
    23  	return &Builder{
    24  		jumpAddr:         make(map[int]uint32),
    25  		jumpPlaceholders: make(map[int][]int),
    26  	}
    27  }
    28  
    29  // AddInt64 adds a pushdata instruction for an integer value.
    30  func (b *Builder) AddUint64(n uint64) *Builder {
    31  	b.program = append(b.program, vm.PushDataUint64(n)...)
    32  	return b
    33  }
    34  
    35  // AddData adds a pushdata instruction for a given byte string.
    36  func (b *Builder) AddData(data []byte) *Builder {
    37  	b.program = append(b.program, vm.PushDataBytes(data)...)
    38  	return b
    39  }
    40  
    41  // AddRawBytes simply appends the given bytes to the program. (It does
    42  // not introduce a pushdata opcode.)
    43  func (b *Builder) AddRawBytes(data []byte) *Builder {
    44  	b.program = append(b.program, data...)
    45  	return b
    46  }
    47  
    48  // AddOp adds the given opcode to the program.
    49  func (b *Builder) AddOp(op vm.Op) *Builder {
    50  	b.program = append(b.program, byte(op))
    51  	return b
    52  }
    53  
    54  // NewJumpTarget allocates a number that can be used as a jump target
    55  // in AddJump and AddJumpIf. Call SetJumpTarget to associate the
    56  // number with a program location.
    57  func (b *Builder) NewJumpTarget() int {
    58  	b.jumpCounter++
    59  	return b.jumpCounter
    60  }
    61  
    62  // AddJump adds a JUMP opcode whose target is the given target
    63  // number. The actual program location of the target does not need to
    64  // be known yet, as long as SetJumpTarget is called before Build.
    65  func (b *Builder) AddJump(target int) *Builder {
    66  	return b.addJump(vm.OP_JUMP, target)
    67  }
    68  
    69  // AddJump adds a JUMPIF opcode whose target is the given target
    70  // number. The actual program location of the target does not need to
    71  // be known yet, as long as SetJumpTarget is called before Build.
    72  func (b *Builder) AddJumpIf(target int) *Builder {
    73  	return b.addJump(vm.OP_JUMPIF, target)
    74  }
    75  
    76  func (b *Builder) addJump(op vm.Op, target int) *Builder {
    77  	b.AddOp(op)
    78  	b.jumpPlaceholders[target] = append(b.jumpPlaceholders[target], len(b.program))
    79  	b.AddRawBytes([]byte{0, 0, 0, 0})
    80  	return b
    81  }
    82  
    83  // SetJumpTarget associates the given jump-target number with the
    84  // current position in the program - namely, the program's length,
    85  // such that the first instruction executed by a jump using this
    86  // target will be whatever instruction is added next. It is legal for
    87  // SetJumpTarget to be called at the end of the program, causing jumps
    88  // using that target to fall off the end. There must be a call to
    89  // SetJumpTarget for every jump target used before any call to Build.
    90  func (b *Builder) SetJumpTarget(target int) *Builder {
    91  	b.jumpAddr[target] = uint32(len(b.program))
    92  	return b
    93  }
    94  
    95  var ErrUnresolvedJump = errors.New("unresolved jump target")
    96  
    97  // Build produces the bytecode of the program. It first resolves any
    98  // jumps in the program by filling in the addresses of their
    99  // targets. This requires SetJumpTarget to be called prior to Build
   100  // for each jump target used (in a call to AddJump or AddJumpIf). If
   101  // any target's address hasn't been set in this way, this function
   102  // produces ErrUnresolvedJump. There are no other error conditions.
   103  func (b *Builder) Build() ([]byte, error) {
   104  	for target, placeholders := range b.jumpPlaceholders {
   105  		addr, ok := b.jumpAddr[target]
   106  		if !ok {
   107  			return nil, errors.Wrapf(ErrUnresolvedJump, "target %d", target)
   108  		}
   109  		for _, placeholder := range placeholders {
   110  			binary.LittleEndian.PutUint32(b.program[placeholder:placeholder+4], addr)
   111  		}
   112  	}
   113  	return b.program, nil
   114  }