Add inline macros

4 files changed, 133 insertions, 26 deletions

diff --git a/instructions.go b/instructions.go
index 0a7afb9..94b48bd 100644
--- a/instructions.go
+++ b/instructions.go
@@ -8,9 +8,10 @@ import (
 type ParamType func(labels *Labels, defs *Definitions, param string) (uint16, error)
 
 type InstructionParams struct {
-	Types     []ParamType
-	Assembler func(currentAddress uint16, args []uint16) ([]uint8, error)
-	Wildcard  bool
+	Types          []ParamType
+	Assembler      func(currentAddress uint16, args []uint16) ([]uint8, error)
+	Wildcard       bool
+	MacroForbidden bool
 }
 
 type InstructionSet map[string][]InstructionParams
@@ -283,7 +284,26 @@ func InstructionSetNew() InstructionSet {
 				return []byte{0b00100000 | (uint8(args[0]) << 3), uint8(args[1])}, nil
 			},
 		},
-		// TODO: Add the relative thingies somehow
+		{
+			Types: []ParamType{MacroLabel},
+			Assembler: func(currentAddress uint16, args []uint16) ([]byte, error) {
+				relativeAddress, err := absoluteJPValueToRelative(currentAddress, args[0])
+				if err != nil {
+					return nil, err
+				}
+				return []byte{0b00011000, relativeAddress}, nil
+			},
+		},
+		{
+			Types: []ParamType{Condition, MacroLabel},
+			Assembler: func(currentAddress uint16, args []uint16) ([]byte, error) {
+				relativeAddress, err := absoluteJPValueToRelative(currentAddress, args[1])
+				if err != nil {
+					return nil, err
+				}
+				return []byte{0b00100000 | (uint8(args[0]) << 3), relativeAddress}, nil
+			},
+		},
 		{
 			Types: []ParamType{Raw16},
 			Assembler: func(currentAddress uint16, args []uint16) ([]byte, error) {
@@ -293,6 +313,7 @@ func InstructionSetNew() InstructionSet {
 				}
 				return []byte{0b00011000, relativeAddress}, nil
 			},
+			MacroForbidden: true,
 		},
 		{
 			Types: []ParamType{Condition, Raw16},
@@ -303,6 +324,7 @@ func InstructionSetNew() InstructionSet {
 				}
 				return []byte{0b00100000 | (uint8(args[0]) << 3), relativeAddress}, nil
 			},
+			MacroForbidden: true,
 		},
 	}
 	result["CALL"] = []InstructionParams{
@@ -494,6 +516,7 @@ func InstructionSetNew() InstructionSet {
 func (set InstructionSet) Parse(
 	labels *Labels,
 	defs *Definitions,
+	isMacro bool,
 	currentAddress uint16,
 	line string,
 ) ([]byte, error) {
@@ -533,6 +556,10 @@ instruction_param_loop:
 			parsed_params[i] = parsed
 		}
 
+		if instrParam.MacroForbidden && isMacro {
+			return nil, fmt.Errorf("This instruction cannot be used with this set of params inside of a macro")
+		}
+
 		return instrParam.Assembler(currentAddress, parsed_params)
 	}
 	return nil, fmt.Errorf(
diff --git a/macros.go b/macros.go
index 5003524..e08562c 100644
--- a/macros.go
+++ b/macros.go
@@ -10,6 +10,17 @@ import (
 
 var MacroInstructions = NewInstructionSetMacros()
 
+func InlineMacroAssembler(b []byte) []InstructionParams {
+	return []InstructionParams{
+		{
+			Types: []ParamType{},
+			Assembler: func(currentAddress uint16, args []uint16) ([]uint8, error) {
+				return b, nil
+			},
+		},
+	}
+}
+
 func NewInstructionSetMacros() InstructionSet {
 	result := make(InstructionSet)
 
@@ -17,14 +28,9 @@ func NewInstructionSetMacros() InstructionSet {
 		{
 			Types: []ParamType{Raw16},
 			Assembler: func(currentAddress uint16, args []uint16) ([]byte, error) {
-				fmt.Printf(
-					".PADTO 0x%04x, currentAddress: 0x%04x, inserting: 0x%04x\n",
-					args[0],
-					currentAddress,
-					args[0]-currentAddress,
-				)
 				return make([]byte, args[0]-currentAddress), nil
 			},
+			MacroForbidden: true,
 		},
 	}
 
@@ -54,9 +60,10 @@ type (
 
 func MacroParse(
 	line string,
+	lines []string,
 	result *[]byte,
 	state *ProgramState,
-	_ *int, // line_nb
+	line_nb *int, // line_nb
 	is_first_pass bool,
 	offset uint,
 ) error {
@@ -71,6 +78,7 @@ func MacroParse(
 		new_instruction, err := MacroInstructions.Parse(
 			&state.Labels,
 			&state.Defs,
+			state.IsMacro,
 			uint16(uint(len(*result))+offset),
 			line,
 		)
@@ -80,7 +88,7 @@ func MacroParse(
 
 		*result = append(*result, new_instruction...)
 		return nil
-	} else if macroName == ".INCLUDE" {
+	} else if macroName == ".INCLUDE" && !state.IsMacro {
 		filePath := strings.Trim(strings.TrimSpace(strings.TrimPrefix(line, ".INCLUDE")), "\"'")
 
 		input_file, err := os.Open(filePath)
@@ -107,7 +115,7 @@ func MacroParse(
 			}
 			*result = append(*result, included...)
 		}
-	} else if macroName == ".DEFINE" {
+	} else if macroName == ".DEFINE" && !state.IsMacro {
 		if len(words) != 3 {
 			return fmt.Errorf(".DEFINE must have 2 arguments (%v)", words)
 		}
@@ -132,8 +140,56 @@ func MacroParse(
 		}
 
 		state.Defs[name] = definedValue
+	} else if macroName == ".MACRODEF" && !state.IsMacro {
+		if len(words) != 2 {
+			return fmt.Errorf(".MACRODEF should have one argument, followed by the definition")
+		}
+		definedMacroName := strings.ToUpper(words[1])
+		(*line_nb) += 1
+		macroContent := []byte{}
+		for *line_nb < len(lines) && strings.TrimSpace(strings.Split(lines[*line_nb], ";")[0]) != ".END" {
+			macroContent = append(macroContent, (lines[*line_nb] + "\n")...)
+			(*line_nb) += 1
+		}
+
+		state := ProgramState{
+			Labels:  Clone(state.Labels),
+			Defs:    Clone(state.Defs),
+			IsMacro: true,
+		}
+
+		if is_first_pass {
+			if _, ok := MacroInstructions[definedMacroName]; ok {
+				return fmt.Errorf("Macro %s is already defined", definedMacroName)
+			}
+
+			new_instructions, err := firstPass("MACRO$"+definedMacroName, macroContent, 0, &state)
+			if err != nil {
+				return err
+			}
+			MacroInstructions["."+definedMacroName] = InlineMacroAssembler(new_instructions)
+		} else {
+			_, err := firstPass("MACRO$"+definedMacroName, macroContent, 0, &state)
+			if err != nil {
+				return err
+			}
+			new_instructions, err := secondPass("MACRO$"+definedMacroName, macroContent, offset, state)
+			if err != nil {
+				return err
+			}
+			MacroInstructions["."+definedMacroName] = InlineMacroAssembler(new_instructions)
+
+		}
 	} else {
 		return fmt.Errorf("Unknown macro \"%s\"", macroName)
 	}
 	return nil
 }
+
+func Clone[K comparable, V any](arg map[K]V) map[K]V {
+	result := make(map[K]V)
+	for k, v := range arg {
+		result[k] = v
+	}
+	return result
+}
diff --git a/main.go b/main.go
index 30bda30..8100e43 100644
--- a/main.go
+++ b/main.go
@@ -13,14 +13,16 @@ type (
 )
 
 type ProgramState struct {
-	Labels Labels
-	Defs   Definitions
+	Labels  Labels
+	Defs    Definitions
+	IsMacro bool
 }
 
 func parseFile(input_file_name string, input []byte, offset uint) ([]byte, error) {
 	state := ProgramState{
-		Labels: make(map[string]uint),
-		Defs:   make(map[string]any),
+		Labels:  make(map[string]uint),
+		Defs:    make(map[string]any),
+		IsMacro: false,
 	}
 
 	_, err := firstPass(input_file_name, input, offset, &state)
@@ -51,15 +53,21 @@ func firstPass(
 			for _, label := range parts[:len(parts)-1] {
 				label = strings.TrimSpace(strings.ToUpper(label))
 				if _, ok := state.Labels[label]; ok {
-					fmt.Fprintf(
-						os.Stderr,
+					return nil, fmt.Errorf(
 						"File %s, line %d:\nLabel %s is already defined",
 						input_file_name,
 						line_nb,
 						label,
 					)
-					os.Exit(1)
 				}
+
+				if label[0] == '$' && !state.IsMacro {
+					return nil, fmt.Errorf("Labels starting with $ can only be used inside macros")
+				}
+				if label[0] != '$' && state.IsMacro {
+					return nil, fmt.Errorf("Labels inside a macro must start with $")
+				}
+
 				state.Labels[label] = uint(len(result)) + offset
 			}
 
@@ -70,7 +78,7 @@ func firstPass(
 
 		// nil sets all the labels and defintion to 0 & thus, to not crash JR, the currentAddress should also be 0
 		if strings.HasPrefix(line, ".") {
-			err := MacroParse(line, &result, state, &line_nb, true, offset)
+			err := MacroParse(line, lines, &result, state, &line_nb, true, offset)
 			if err != nil {
 				return nil, fmt.Errorf(
 					"File %s, line %d (1st pass|macro):\n%w",
@@ -80,7 +88,7 @@ func firstPass(
 				)
 			}
 		} else {
-			next_instruction, err := Instructions.Parse(nil, &state.Defs, 0, line)
+			next_instruction, err := Instructions.Parse(nil, &state.Defs, state.IsMacro, 0, line)
 			if err != nil {
 				return nil, fmt.Errorf(
 					"File %s, line %d (1st pass):\n%w",
@@ -123,7 +131,7 @@ func secondPass(
 		line = strings.TrimSpace(line)
 
 		if strings.HasPrefix(line, ".") {
-			err := MacroParse(line, &result, &state, &line_nb, false, offset)
+			err := MacroParse(line, lines, &result, &state, &line_nb, false, offset)
 			if err != nil {
 				return nil, fmt.Errorf(
 					"File %s, line %d (2nd pass|macro):\n%w",
@@ -133,7 +141,7 @@ func secondPass(
 				)
 			}
 		} else {
-			next_instruction, err := Instructions.Parse(&state.Labels, &state.Defs, uint16(uint(len(result))+offset), line)
+			next_instruction, err := Instructions.Parse(&state.Labels, &state.Defs, state.IsMacro, uint16(uint(len(result))+offset), line)
 			if err != nil {
 				return nil, fmt.Errorf(
 					"File %s, line %d (2nd pass): %w",
diff --git a/parameters.go b/parameters.go
index 8a3174b..5e234f5 100644
--- a/parameters.go
+++ b/parameters.go
@@ -120,7 +120,7 @@ func Raw16(labels *Labels, defs *Definitions, param string) (uint16, error) {
 		return uint16(res), nil
 	}
 
-	if strings.HasPrefix(param, "=") {
+	if strings.HasPrefix(param, "=") && !strings.HasPrefix(param, "=$") {
 		if labels == nil {
 			return 0, nil
 		}
@@ -144,6 +144,23 @@ func Raw16(labels *Labels, defs *Definitions, param string) (uint16, error) {
 	return uint16(res), err
 }
 
+func MacroLabel(labels *Labels, defs *Definitions, param string) (uint16, error) {
+	if strings.HasPrefix(param, "=$") {
+		if labels == nil {
+			return 0, nil
+		}
+
+		label := strings.ToUpper(strings.TrimPrefix(param, "="))
+		labelValue, ok := (*labels)[label]
+		if !ok {
+			return 0, fmt.Errorf("Label \"%s\" not found", label)
+		}
+
+		return uint16(labelValue), nil
+	}
+	return 0, fmt.Errorf("Invalid macro label")
+}
+
 func Reg16Indirect(_ *Labels, _ *Definitions, param string) (uint16, error) {
 	switch param {
 	case "(BC)":
@@ -202,7 +219,6 @@ func Raw16Indirect(labels *Labels, defs *Definitions, param string) (uint16, err
 		if !ok {
 			return 0, fmt.Errorf("$%s is of type %T but Indirect16b expected", param, res)
 		}
-		fmt.Println("YOU CAN LOOSE YOUR MIND NOW ! IT OFFICIALLY DOESN'T MAKE SENSE", param, res)
 		return uint16(res), nil
 	}
 	if len(param) < 2 || param[0] != '(' || param[len(param)-1] != ')' {