github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/vm/stack.go (about)

     1  package vm
     2  
     3  import (
     4  	"encoding/json"
     5  	"errors"
     6  	"fmt"
     7  	"math/big"
     8  
     9  	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
    10  )
    11  
    12  // Stack implementation for the neo-go virtual machine. The stack with its LIFO
    13  // semantics is emulated from a simple slice, where the top of the stack corresponds
    14  // to the latest element of this slice. Pushes are appends to this slice, pops are
    15  // slice resizes.
    16  
    17  // Element represents an element on the stack. Technically, it's a wrapper around
    18  // stackitem.Item interface to provide some API simplification for the VM.
    19  type Element struct {
    20  	value stackitem.Item
    21  }
    22  
    23  // NewElement returns a new Element object, with its underlying value inferred
    24  // to the corresponding type.
    25  func NewElement(v any) Element {
    26  	return Element{stackitem.Make(v)}
    27  }
    28  
    29  // Item returns the Item contained in the element.
    30  func (e Element) Item() stackitem.Item {
    31  	return e.value
    32  }
    33  
    34  // Value returns the value of the Item contained in the element.
    35  func (e Element) Value() any {
    36  	return e.value.Value()
    37  }
    38  
    39  // BigInt attempts to get the underlying value of the element as a big integer.
    40  // It will panic if the assertion has failed, which will be caught by the VM.
    41  func (e Element) BigInt() *big.Int {
    42  	val, err := e.value.TryInteger()
    43  	if err != nil {
    44  		panic(err)
    45  	}
    46  	return val
    47  }
    48  
    49  // Bool converts the underlying value of the element to a boolean if it's
    50  // possible to do so. Otherwise, it will panic.
    51  func (e Element) Bool() bool {
    52  	b, err := e.value.TryBool()
    53  	if err != nil {
    54  		panic(err)
    55  	}
    56  	return b
    57  }
    58  
    59  // Bytes attempts to get the underlying value of the element as a byte array.
    60  // It will panic if the assertion has failed, which will be caught by the VM.
    61  func (e Element) Bytes() []byte {
    62  	bs, err := e.value.TryBytes()
    63  	if err != nil {
    64  		panic(err)
    65  	}
    66  	return bs
    67  }
    68  
    69  // BytesOrNil attempts to get the underlying value of the element as a byte array or nil.
    70  // It will panic if the assertion has failed, which will be caught by the VM.
    71  func (e Element) BytesOrNil() []byte {
    72  	if _, ok := e.value.(stackitem.Null); ok {
    73  		return nil
    74  	}
    75  	bs, err := e.value.TryBytes()
    76  	if err != nil {
    77  		panic(err)
    78  	}
    79  	return bs
    80  }
    81  
    82  // String attempts to get a string from the element value.
    83  // It is assumed to be used in interops and panics if the string is not a valid UTF-8 byte sequence.
    84  func (e Element) String() string {
    85  	s, err := stackitem.ToString(e.value)
    86  	if err != nil {
    87  		panic(err)
    88  	}
    89  	return s
    90  }
    91  
    92  // Array attempts to get the underlying value of the element as an array of
    93  // other items. It will panic if the item type is different, which will be caught
    94  // by the VM.
    95  func (e Element) Array() []stackitem.Item {
    96  	switch t := e.value.(type) {
    97  	case *stackitem.Array:
    98  		return t.Value().([]stackitem.Item)
    99  	case *stackitem.Struct:
   100  		return t.Value().([]stackitem.Item)
   101  	default:
   102  		panic("element is not an array")
   103  	}
   104  }
   105  
   106  // Interop attempts to get the underlying value of the element
   107  // as an interop item.
   108  func (e Element) Interop() *stackitem.Interop {
   109  	switch t := e.value.(type) {
   110  	case *stackitem.Interop:
   111  		return t
   112  	default:
   113  		panic("element is not an interop")
   114  	}
   115  }
   116  
   117  // Stack represents a Stack backed by a slice of Elements.
   118  type Stack struct {
   119  	elems []Element
   120  	name  string
   121  	refs  *refCounter
   122  }
   123  
   124  // NewStack returns a new stack name by the given name.
   125  func NewStack(n string) *Stack {
   126  	return newStack(n, newRefCounter())
   127  }
   128  
   129  func newStack(n string, refc *refCounter) *Stack {
   130  	s := new(Stack)
   131  	s.elems = make([]Element, 0, 16) // Most of uses are expected to fit into 16 elements.
   132  	initStack(s, n, refc)
   133  	return s
   134  }
   135  
   136  func subStack(old *Stack) *Stack {
   137  	s := new(Stack)
   138  	*s = *old
   139  	s.elems = s.elems[len(s.elems):]
   140  	return s
   141  }
   142  
   143  func initStack(s *Stack, n string, refc *refCounter) {
   144  	s.name = n
   145  	s.refs = refc
   146  	s.Clear()
   147  }
   148  
   149  // Clear clears all elements on the stack and set its length to 0.
   150  func (s *Stack) Clear() {
   151  	if s.elems != nil {
   152  		for _, el := range s.elems {
   153  			s.refs.Remove(el.value)
   154  		}
   155  		s.elems = s.elems[:0]
   156  	}
   157  }
   158  
   159  // Len returns the number of elements that are on the stack.
   160  func (s *Stack) Len() int {
   161  	return len(s.elems)
   162  }
   163  
   164  // InsertAt inserts the given item (n) deep on the stack.
   165  // Be very careful using it and _always_ check n before invocation
   166  // as it will panic otherwise.
   167  func (s *Stack) InsertAt(e Element, n int) {
   168  	l := len(s.elems)
   169  	s.elems = append(s.elems, e)
   170  	copy(s.elems[l-n+1:], s.elems[l-n:l])
   171  	s.elems[l-n] = e
   172  	s.refs.Add(e.value)
   173  }
   174  
   175  // Push pushes the given element on the stack.
   176  func (s *Stack) Push(e Element) {
   177  	s.elems = append(s.elems, e)
   178  	s.refs.Add(e.value)
   179  }
   180  
   181  // PushItem pushes an Item to the stack.
   182  func (s *Stack) PushItem(i stackitem.Item) {
   183  	s.Push(Element{i})
   184  }
   185  
   186  // PushVal pushes the given value on the stack. It will infer the
   187  // underlying Item to its corresponding type.
   188  func (s *Stack) PushVal(v any) {
   189  	s.Push(NewElement(v))
   190  }
   191  
   192  // Pop removes and returns the element on top of the stack. It panics if the stack is
   193  // empty.
   194  func (s *Stack) Pop() Element {
   195  	l := len(s.elems)
   196  	e := s.elems[l-1]
   197  	s.elems = s.elems[:l-1]
   198  	s.refs.Remove(e.value)
   199  	return e
   200  }
   201  
   202  // Top returns the element on top of the stack. Nil if the stack
   203  // is empty.
   204  func (s *Stack) Top() Element {
   205  	if len(s.elems) == 0 {
   206  		return Element{}
   207  	}
   208  	return s.elems[len(s.elems)-1]
   209  }
   210  
   211  // Back returns the element at the end of the stack. Nil if the stack
   212  // is empty.
   213  func (s *Stack) Back() Element {
   214  	if len(s.elems) == 0 {
   215  		return Element{}
   216  	}
   217  	return s.elems[0]
   218  }
   219  
   220  // Peek returns the element (n) far in the stack beginning from
   221  // the top of the stack. For n == 0 it's, effectively, the same as Top,
   222  // but it'll panic if the stack is empty.
   223  func (s *Stack) Peek(n int) Element {
   224  	n = len(s.elems) - n - 1
   225  	return s.elems[n]
   226  }
   227  
   228  // RemoveAt removes the element (n) deep on the stack beginning
   229  // from the top of the stack. It panics if called with out of bounds n.
   230  func (s *Stack) RemoveAt(n int) Element {
   231  	l := len(s.elems)
   232  	e := s.elems[l-1-n]
   233  	s.elems = append(s.elems[:l-1-n], s.elems[l-n:]...)
   234  	s.refs.Remove(e.value)
   235  	return e
   236  }
   237  
   238  // Dup duplicates and returns the element at position n.
   239  // Dup is used for copying elements on the top of its own stack.
   240  //
   241  //	s.Push(s.Peek(0)) // will result in unexpected behavior.
   242  //	s.Push(s.Dup(0)) // is the correct approach.
   243  func (s *Stack) Dup(n int) Element {
   244  	e := s.Peek(n)
   245  	return Element{e.value.Dup()}
   246  }
   247  
   248  // Iter iterates over all elements int the stack, starting from the top
   249  // of the stack.
   250  //
   251  //	s.Iter(func(elem *Element) {
   252  //		// do something with the element.
   253  //	})
   254  func (s *Stack) Iter(f func(Element)) {
   255  	for i := len(s.elems) - 1; i >= 0; i-- {
   256  		f(s.elems[i])
   257  	}
   258  }
   259  
   260  // IterBack iterates over all elements of the stack, starting from the bottom
   261  // of the stack.
   262  //
   263  //	s.IterBack(func(elem *Element) {
   264  //		// do something with the element.
   265  //	})
   266  func (s *Stack) IterBack(f func(Element)) {
   267  	for i := 0; i < len(s.elems); i++ {
   268  		f(s.elems[i])
   269  	}
   270  }
   271  
   272  // Swap swaps two elements on the stack without popping and pushing them.
   273  func (s *Stack) Swap(n1, n2 int) error {
   274  	if n1 < 0 || n2 < 0 {
   275  		return errors.New("negative index")
   276  	}
   277  	l := len(s.elems)
   278  	if n1 >= l || n2 >= l {
   279  		return errors.New("too big index")
   280  	}
   281  	s.elems[l-n1-1], s.elems[l-n2-1] = s.elems[l-n2-1], s.elems[l-n1-1]
   282  	return nil
   283  }
   284  
   285  // ReverseTop reverses top n items of the stack.
   286  func (s *Stack) ReverseTop(n int) error {
   287  	l := len(s.elems)
   288  	if n < 0 {
   289  		return errors.New("negative index")
   290  	} else if n > l {
   291  		return errors.New("too big index")
   292  	} else if n <= 1 {
   293  		return nil
   294  	}
   295  
   296  	for i, j := l-n, l-1; i <= j; i, j = i+1, j-1 {
   297  		s.elems[i], s.elems[j] = s.elems[j], s.elems[i]
   298  	}
   299  	return nil
   300  }
   301  
   302  // Roll brings an item with the given index to the top of the stack moving all
   303  // other elements down accordingly. It does all of that without popping and
   304  // pushing elements.
   305  func (s *Stack) Roll(n int) error {
   306  	if n < 0 {
   307  		return errors.New("negative index")
   308  	}
   309  	l := len(s.elems)
   310  	if n >= l {
   311  		return errors.New("too big index")
   312  	}
   313  	if n == 0 {
   314  		return nil
   315  	}
   316  	e := s.elems[l-1-n]
   317  	copy(s.elems[l-1-n:], s.elems[l-n:])
   318  	s.elems[l-1] = e
   319  	return nil
   320  }
   321  
   322  // PopSigElements pops keys or signatures from the stack as needed for
   323  // CHECKMULTISIG.
   324  func (s *Stack) PopSigElements() ([][]byte, error) {
   325  	var num int
   326  	var elems [][]byte
   327  	if s.Len() == 0 {
   328  		return nil, fmt.Errorf("nothing on the stack")
   329  	}
   330  	item := s.Pop()
   331  	switch item.value.(type) {
   332  	case *stackitem.Array:
   333  		num = len(item.Array())
   334  		if num < 1 {
   335  			return nil, fmt.Errorf("less than one element in the array")
   336  		}
   337  		elems = make([][]byte, num)
   338  		for k, v := range item.Array() {
   339  			b, ok := v.Value().([]byte)
   340  			if !ok {
   341  				return nil, fmt.Errorf("bad element %s", v.String())
   342  			}
   343  			elems[k] = b
   344  		}
   345  	default:
   346  		num = int(item.BigInt().Int64())
   347  		if num < 1 || num > s.Len() {
   348  			return nil, fmt.Errorf("wrong number of elements: need %d, have %d", num, s.Len())
   349  		}
   350  		elems = make([][]byte, num)
   351  		for i := 0; i < num; i++ {
   352  			elems[i] = s.Pop().Bytes()
   353  		}
   354  	}
   355  	return elems, nil
   356  }
   357  
   358  // ToArray converts the stack to an array of stackitems with the top item being the last.
   359  func (s *Stack) ToArray() []stackitem.Item {
   360  	items := make([]stackitem.Item, 0, len(s.elems))
   361  	s.IterBack(func(e Element) {
   362  		items = append(items, e.Item())
   363  	})
   364  	return items
   365  }
   366  
   367  // MarshalJSON implements the JSON marshalling interface.
   368  func (s *Stack) MarshalJSON() ([]byte, error) {
   369  	items := s.ToArray()
   370  	arr := make([]json.RawMessage, len(items))
   371  	for i := range items {
   372  		data, err := stackitem.ToJSONWithTypes(items[i])
   373  		if err == nil {
   374  			arr[i] = data
   375  		}
   376  	}
   377  	return json.Marshal(arr)
   378  }