add REPT support (#11)

assembler/address_assigning_step.go

@@ -177,7 +177,7 @@ func assignEnumAddress(aa *addressAssign, e ast.Enum) (uint64, error) {
 
 func assignEnumEndAddress(aa *addressAssign) (uint64, error) {
 	if !aa.enumActive {
-		return 0, errors.New("enum outside of enum context")
+		return 0, errors.New("enum end outside of enum context")
 	}
 
 	aa.enumActive = false

assembler/assembler_asm6_test.go

@@ -366,6 +366,29 @@ func TestAssemblerAsm6Enum(t *testing.T) {
 	assert.Equal(t, expected, b)
 }
 
+var asm6ReptCode = `
+.segment "HEADER"
+
+i=0
+REPT 3
+	DB i
+	i=i+1
+ENDR
+DB 0xff
+`
+
+func TestAssemblerAsm6Rept(t *testing.T) {
+	b, err := runAsm6Test(t, unitTestConfig, asm6ReptCode)
+	assert.NoError(t, err)
+	expected := []byte{
+		0x00, // 1 item
+		0x01, // 1 item
+		0x02, // 1 item
+		0xff, // 1 item
+	}
+	assert.Equal(t, expected, b)
+}
+
 func runAsm6Test(t *testing.T, testConfig, testCode string) ([]byte, error) {
 	t.Helper()
 

assembler/expression_evaluation_step.go

@@ -39,12 +39,14 @@ func evaluateExpressionsStep(asm *Assembler) error {
 	for segNr, seg := range asm.segmentsOrder {
 		nodes := make([]ast.Node, 0, len(seg.nodes))
 
-		for nodeNr, node := range seg.nodes {
-			remove, err := evaluateNode(&expEval, node)
+		// nolint:intrange // seg.nodes gets modified in the loop
+		for nodeNr := 0; nodeNr < len(seg.nodes); nodeNr++ {
+			node := seg.nodes[nodeNr]
+			removeNode, err := evaluateNode(&expEval, seg, nodeNr, node)
 			if err != nil {
 				return fmt.Errorf("evaluating node %d in segment %d: %w", nodeNr, segNr, err)
 			}
-			if !remove {
+			if !removeNode {
 				nodes = append(nodes, node)
 			}
 		}
@@ -61,8 +63,8 @@ func evaluateExpressionsStep(asm *Assembler) error {
 // evaluateNode evaluates a node and returns whether the node should be removed.
 // This is useful for conditional nodes with an expression that does not match and
 // that wraps other nodes.
-// nolint:cyclop
-func evaluateNode(expEval *expressionEvaluation, node any) (bool, error) {
+// nolint:cyclop,funlen
+func evaluateNode(expEval *expressionEvaluation, seg *segment, currentNodeIndex int, node any) (bool, error) {
 	// always handle conditional nodes
 	switch n := node.(type) {
 	case ast.If:
@@ -112,6 +114,15 @@ func evaluateNode(expEval *expressionEvaluation, node any) (bool, error) {
 			return false, fmt.Errorf("evaluating enum expression: %w", err)
 		}
 
+	case ast.Rept:
+		if err := parseRept(expEval, n, seg, currentNodeIndex); err != nil {
+			return false, err
+		}
+		return true, nil
+
+	case ast.Endr:
+		return true, nil
+
 	case *data:
 		return false, parseDataExpression(expEval, n)
 
@@ -278,3 +289,58 @@ func parseElseIfCondition(expEval *expressionEvaluation, cond ast.ElseIf) error
 	expEval.currentContext.processNodes = conditionMet
 	return nil
 }
+
+func parseRept(expEval *expressionEvaluation, rept ast.Rept, seg *segment, currentNodeIndex int) error {
+	if rept.Count.IsEvaluatedAtAddressAssign() {
+		return errExpressionCantReferenceProgramCounter
+	}
+
+	_, err := rept.Count.Evaluate(expEval.currentScope, expEval.arch.AddressWidth)
+	if err != nil {
+		return fmt.Errorf("evaluating if condition at program counter: %w", err)
+	}
+
+	count, err := rept.Count.IntValue()
+	if err != nil {
+		return fmt.Errorf("getting rept count: %w", err)
+	}
+	if count <= 0 {
+		return errors.New("rept count must be positive")
+	}
+
+	var nodes []ast.Node
+	var reptEnded bool
+
+	for i := currentNodeIndex + 1; i < len(seg.nodes); i++ {
+		node := seg.nodes[i]
+		if _, ok := node.(ast.Endr); !ok {
+			nodes = append(nodes, node)
+		} else {
+			reptEnded = true
+			break
+		}
+	}
+
+	if !reptEnded {
+		return errors.New("rept without endr found")
+	}
+
+	// insert the nodes count-1 times, as the first insertion are the existing nodes
+	count--
+	nodesToInsert := make([]ast.Node, 0, len(nodes)*int(count))
+	for range count {
+		for _, node := range nodes {
+			nodesToInsert = append(nodesToInsert, node.Copy())
+		}
+	}
+
+	// copy nodes up to endr
+	nodes = seg.nodes[:currentNodeIndex+len(nodesToInsert)-1]
+	// append now node copies
+	nodes = append(nodes, nodesToInsert...)
+	// append nodes after endr
+	nodes = append(nodes, seg.nodes[currentNodeIndex+len(nodesToInsert):]...)
+
+	seg.nodes = nodes
+	return nil
+}

assembler/nodes.go

@@ -144,7 +144,7 @@ type symbol struct {
 // Copy returns a copy of the symbol node.
 func (s *symbol) Copy() ast.Node {
 	return &symbol{
-		Symbol: s.Symbol,
+		Symbol: s.Symbol.Copy(),
 	}
 }
 

assembler/process_macros_step.go

@@ -8,12 +8,14 @@ import (
 	"github.com/retroenv/retroasm/parser/ast"
 )
 
-// processMacrosStep processes macro usages and replace them by the macro nodes.
+// processMacrosStep processes macro and rept nodes and replace them by their resolved nodes.
 func processMacrosStep(asm *Assembler) error {
 	for i, seg := range asm.segmentsOrder {
 		segmentNodesResolved := make([]ast.Node, 0, len(seg.nodes))
 
-		for _, node := range seg.nodes {
+		for j := range seg.nodes {
+			node := seg.nodes[j]
+
 			switch n := node.(type) {
 			case ast.Identifier:
 				nodes, err := resolveMacroUsage(asm, n)

expression/expression.go

@@ -58,6 +58,12 @@ func (e *Expression) Copy() *Expression {
 	}
 }
 
+// CopyExpression creates a copy of the expression.
+// Secondary copy function to avoid cyclic dependency.
+func (e *Expression) CopyExpression() any {
+	return e.Copy()
+}
+
 // SetEvaluateOnce sets the evaluate once flag for the expression.
 // = will declare an alias that is evaluated on processing of the node,
 // EQU will declare an expression that is evaluated on every usage and

parser/ast/rept.go

@@ -0,0 +1,48 @@
+package ast
+
+import (
+	"github.com/retroenv/retroasm/expression"
+	"github.com/retroenv/retroasm/lexer/token"
+)
+
+// Rept ...
+type Rept struct {
+	*node
+
+	Count *expression.Expression
+}
+
+// NewRept returns a new rept node.
+func NewRept(count []token.Token) Rept {
+	return Rept{
+		node:  &node{},
+		Count: expression.New(count...),
+	}
+}
+
+// Copy returns a copy of the rept node.
+func (r Rept) Copy() Node {
+	return Rept{
+		node:  r.node,
+		Count: r.Count.Copy(),
+	}
+}
+
+// Endr ...
+type Endr struct {
+	*node
+}
+
+// NewEndr returns a new rept end node.
+func NewEndr() Endr {
+	return Endr{
+		node: &node{},
+	}
+}
+
+// Copy returns a copy of the rept end node.
+func (e Endr) Copy() Node {
+	return Endr{
+		node: e.node,
+	}
+}

parser/directives/directives.go

@@ -46,6 +46,7 @@ var Handlers = map[string]Handler{
 	"endif":      Endif,  // asm6
 	"ende":       Ende,   // asm6
 	"endproc":    EndProc,
+	"endr":       Endr,      // asm6
 	"enum":       Enum,      // asm6
 	"error":      Error,     // asm6
 	"fillvalue":  FillValue, // asm6
@@ -66,6 +67,7 @@ var Handlers = map[string]Handler{
 	"org":        Base,    // asm6
 	"pad":        Padding, // asm6
 	"proc":       Proc,
+	"rept":       Rept, // asm6
 	"res":        Res,
 	"rsset":      NesasmOffsetCounter,
 	"segment":    Segment,

parser/directives/rept.go

@@ -0,0 +1,24 @@
+package directives
+
+import (
+	"fmt"
+
+	"github.com/retroenv/retroasm/parser/ast"
+)
+
+// Rept ...
+func Rept(p Parser) (ast.Node, error) {
+	p.AdvanceReadPosition(1)
+	countTokens, err := readDataTokens(p, true)
+	if err != nil {
+		return nil, fmt.Errorf("reading rept count tokens: %w", err)
+	}
+
+	return ast.NewRept(countTokens), nil
+}
+
+// Endr ...
+func Endr(p Parser) (ast.Node, error) {
+	p.AdvanceReadPosition(1)
+	return ast.NewEndr(), nil
+}

scope/symbol.go

@@ -20,6 +20,7 @@ const (
 
 // Expression defines the used expression functions.
 type Expression interface {
+	CopyExpression() any
 	Evaluate(scope *Scope, dataWidth int) (any, error)
 	EvaluateAtProgramCounter(scope *Scope, dataWidth int, programCounter uint64) (any, error)
 	IsEvaluatedAtAddressAssign() bool
@@ -48,6 +49,16 @@ func NewSymbol(scope *Scope, name string, typ SymbolType) (*Symbol, error) {
 	return sym, nil
 }
 
+// Copy returns a copy of the symbol.
+func (sym *Symbol) Copy() *Symbol {
+	return &Symbol{
+		name:       sym.name,
+		address:    sym.address,
+		typ:        sym.typ,
+		expression: sym.expression.CopyExpression().(Expression),
+	}
+}
+
 // SetAddress sets the address of the symbol. This is only useful for symbols of type label that
 // gets referenced in code.
 func (sym *Symbol) SetAddress(address uint64) {