1 files changed, 27 insertions, 13 deletions

diff --git a/tools/assemble.c b/tools/assemble.c
index 11bb7c9..b8bd75e 100644
--- a/tools/assemble.c
+++ b/tools/assemble.c
@@ -55,7 +55,8 @@ static uint8_t estimate_instruction_size(struct instruction *instruction)
 	case DATA_NONE :
 		return 2;
 	case DATA_KNOWN :
-	case DATA_UNKNOWN :
+	case DATA_INSTR_ADDR :
+	case DATA_ADDR_AFTER :
 		return 6;
 	case DATA_KNOWN_21_BITS :
 		return 4;
@@ -81,16 +82,26 @@ static void assign_addresses_and_sizes(struct instruction *expr,
 				       uint32_t address)
 {
 	struct instruction *tmp = expr;
+	uint32_t target_addr;
 
 	do {
+		if ((tmp->data.info == DATA_INSTR_ADDR ||
+		     tmp->data.info == DATA_ADDR_AFTER) &&
+		    *tmp->data.data.ptr &&
+		    (*tmp->data.data.ptr)->address_assigned) {
+			target_addr = (*tmp->data.data.ptr)->address;
+
+			if (tmp->data.info == DATA_ADDR_AFTER) {
+				target_addr += estimate_instruction_size
+					(*tmp->data.data.ptr);
+			}
+
+			tmp->data = im(target_addr);
+		}
+
 		tmp->address = address;
 		tmp->address_assigned = true;
 
-		if (tmp->data.info == DATA_UNKNOWN &&
-		    *tmp->data.data.ptr &&
-		    (*tmp->data.data.ptr)->address_assigned)
-			tmp->data = im((*tmp->data.data.ptr)->address);
-
 		address += estimate_instruction_size(tmp);
 		tmp = tmp->next;
 	} while (tmp != expr);
@@ -108,21 +119,24 @@ static void encode_instruction(struct instruction *instruction,
 	uint16_t encoding = instruction->encoding;
 	uint16_t *dest = memory + instruction->address / 2;
 
-	if (instruction->data.info != DATA_NONE) {
-		if (instruction->data.info == DATA_UNKNOWN)
-			im = (*instruction->data.data.ptr)->address;
-		else
-			im = instruction->data.data.im;
+	if (instruction->data.info == DATA_INSTR_ADDR) {
+		im = (*instruction->data.data.ptr)->address;
+	} else if (instruction->data.info == DATA_ADDR_AFTER) {
+		im = (*instruction->data.data.ptr)->address +
+			estimate_instruction_size(*instruction->data.data.ptr);
+	} else if (instruction->data.info != DATA_NONE) {
+		im = instruction->data.data.im;
 	}
 
 	switch (instruction->data.info) {
-	case DATA_UNKNOWN :
+	case DATA_INSTR_ADDR :
+	case DATA_ADDR_AFTER :
 	case DATA_KNOWN :
 		*(dest++) = im_instruction(im >> 22);
 	case DATA_KNOWN_21_BITS :
 		*(dest++) = im_instruction(im >> 7);
 	case DATA_KNOWN_6_BITS :
-		encoding |= (im & 0x7F);
+		encoding |= im & 0x7F;
 	}
 
 	*dest = encoding;