asm: Support add x, y, z<<1 style syntax for the arm64 target

src/asm_tables.cpp

@@ -33,6 +33,7 @@ enum AsmOperandKind : u8 {
 	AsmOperand_Label,
 	AsmOperand_RegisterGroup,
 	AsmOperand_Lane,
+	AsmOperand_RegisterShift,
 
 	AsmOperand_COUNT
 };
@@ -46,6 +47,7 @@ gb_global String const asm_operand_kind_strings[AsmOperand_COUNT] = {
 	str_lit("label"),
 	str_lit("register group"),
 	str_lit("lane"),
+	str_lit("register shift"),
 };
 
 gb_global String const asm_operand_kind_expected_strings[AsmOperand_COUNT] = {
@@ -57,6 +59,7 @@ gb_global String const asm_operand_kind_expected_strings[AsmOperand_COUNT] = {
 	str_lit("a label"),
 	str_lit("a register group"),
 	str_lit("a lane"),
+	str_lit("a register shift"),
 };
 
 enum AsmOperandConstraintKind : u32 {

src/asm_tables_arm64.cpp

@@ -867,8 +867,6 @@ struct Asm_arm64 {
 		// Integer GPR / GPR-or-SP, including shifted- and extended-register forms.
 		case OP_W_REG:      case OP_X_REG:
 		case OP_WSP_REG:    case OP_XSP_REG:
-		case OP_W_SHIFTED:  case OP_X_SHIFTED:
-		case OP_W_EXTENDED: case OP_X_EXTENDED:
 		// SIMD&FP scalar views.
 		case OP_B_REG: case OP_H_REG: case OP_S_REG: case OP_D_REG: case OP_Q_REG:
 		// NEON vector: plain, arrangement, FP16, element-indexed.
@@ -885,6 +883,13 @@ struct Asm_arm64 {
 		case OP_SYS_REG:     // MRS/MSR system-register name -> 16-bit field (cf. riscv CSR)
 			return AsmOperand_Register;
 
+		// ---- Shifted / extended register operands (`reg, lsl #n` / `reg, uxtw #n`) ----
+		// A plain register also fills these slots (shift #0); a shifted register does
+		// NOT fill a plain-register slot. That asymmetry lives in asm_operand_kind_fits.
+		case OP_W_SHIFTED:  case OP_X_SHIFTED:
+		case OP_W_EXTENDED: case OP_X_EXTENDED:
+			return AsmOperand_RegisterShift;
+
 		// ---- Immediates (numeric literals and immediate-encoded selectors) ----
 		case OP_IMM_2:  case OP_IMM_3:  case OP_IMM_4:  case OP_IMM_5:
 		case OP_IMM_6:  case OP_IMM_8:  case OP_IMM_12: case OP_IMM_16:

src/check_asm.cpp

@@ -93,10 +93,28 @@ gb_internal AsmOperandKind determine_asm_operand_kind(Operand const *operand) {
 	case_ast_node(ie, IndexExpr, expr);
 		return AsmOperand_Lane;
 	case_end;
+
+	case_ast_node(be, BinaryExpr, expr);
+		return AsmOperand_RegisterShift;
+	case_end;
+
 	}
 	return AsmOperand_Invalid;
 }
 
+gb_internal bool asm_operand_kind_fits(AsmOperandKind dst, AsmOperandKind src) {
+	if (dst == src) {
+		return true;
+	}
+	switch (dst) {
+	case AsmOperand_Register_Or_Memory:
+		return src == AsmOperand_Register || src == AsmOperand_Memory;
+	case AsmOperand_RegisterShift:
+		return src == AsmOperand_Register;
+	}
+	return false;
+}
+
 gb_internal bool asm_reg_class_compatible(AsmRegClass want, AsmRegClass got) {
 	switch (want) {
 	case AsmRegClass_Integer:
@@ -175,6 +193,15 @@ gb_internal void check_asm_collect_refs(AsmCtx *asm_ctx, PtrSet<Entity *> *refs,
 		check_asm_collect_refs(asm_ctx, refs, expr->IndexExpr.expr,  touched_regs_);
 		check_asm_collect_refs(asm_ctx, refs, expr->IndexExpr.index, touched_regs_);
 		return;
+
+	case Ast_UnaryExpr:
+		check_asm_collect_refs(asm_ctx, refs, expr->UnaryExpr.expr,  touched_regs_);
+		return;
+
+	case Ast_BinaryExpr:
+		check_asm_collect_refs(asm_ctx, refs, expr->BinaryExpr.left,  touched_regs_);
+		check_asm_collect_refs(asm_ctx, refs, expr->BinaryExpr.right, touched_regs_);
+		return;
 	}
 }
 enum AsmMismatch : u8 {
@@ -1484,8 +1511,7 @@ gb_internal bool check_mnemonic(AsmCtx *asm_ctx, CheckerContext *ctx, Entity *tm
 			AsmOperandKind dst = asm_ctx->kind_from_operand_type(type);
 			AsmOperandKind src = determine_asm_operand_kind(operand);
 
-			bool kind_ok = (dst == src) ||
-			               (dst == AsmOperand_Register_Or_Memory && (src == AsmOperand_Register || src == AsmOperand_Memory));
+			bool kind_ok = asm_operand_kind_fits(dst, src);
 
 			// Bias toward wider register slots so an r64 form outranks an otherwise-equal r32 form.
 			width_pref += cast(int)asm_ctx->operand_type_bit_width(type);
@@ -1976,9 +2002,7 @@ gb_internal bool check_mnemonic(AsmCtx *asm_ctx, CheckerContext *ctx, Entity *tm
 			possible_kinds      [i] = dst;
 			possible_class_kinds[i] = asm_ctx->reg_class_from_operand_type(type);
 
-			bool kind_ok = (dst == src) ||
-			               (dst == AsmOperand_Register_Or_Memory
-			                && (src == AsmOperand_Register || src == AsmOperand_Memory));
+			bool kind_ok = asm_operand_kind_fits(dst, src);
 			if (!kind_ok) {
 				valid_spots[i] = false;
 			} else {
@@ -2616,7 +2640,7 @@ gb_internal void check_asm_instruction_operand(AsmCtx *asm_ctx, CheckerContext *
 
 		Operand lhs = {};
 		Operand rhs = {};
-		check_asm_instruction_operand(asm_ctx, ctx, entity, &lhs, ie->expr, false);
+		check_asm_instruction_operand(asm_ctx, ctx, entity, &lhs, ie->expr,  false);
 		check_asm_instruction_operand(asm_ctx, ctx, entity, &rhs, ie->index, false);
 
 		auto lhs_kind = determine_asm_operand_kind(&lhs);
@@ -2666,6 +2690,65 @@ gb_internal void check_asm_instruction_operand(AsmCtx *asm_ctx, CheckerContext *
 		}
 		return;
 	case_end;
+
+	case_ast_node(be, BinaryExpr, expr);
+		switch (be->op.kind) {
+		case Token_Shl:
+		case Token_Shr:
+		case Token_Mul:
+			break;
+		default:
+			error(be->op, "Unsupported operator '%.*s' in an asm operand; only a register shift/scale ('<<', '>>', '*') is allowed here", LIT(be->op.string));
+			return;
+		}
+
+		if (build_context.metrics.arch != TargetArch_arm64) {
+			error(expr, "Asm shifted/scaled register operands are not supported by the target platform");
+			return;
+		}
+
+		Operand reg    = {};
+		Operand amount = {};
+		check_asm_instruction_operand(asm_ctx, ctx, entity, &reg,    be->left,  false);
+		check_asm_instruction_operand(asm_ctx, ctx, entity, &amount, be->right, false);
+		if (reg.mode == Addressing_Invalid || amount.mode == Addressing_Invalid) {
+			return;
+		}
+		if (determine_asm_operand_kind(&reg) != AsmOperand_Register) {
+			gbString s = expr_to_string(reg.expr);
+			error(reg.expr, "The left-hand side of a register shift/scale must be a register, got %s", s);
+			gb_string_free(s);
+			return;
+		}
+
+		if (determine_asm_operand_kind(&amount) != AsmOperand_Immediate) {
+			error(amount.expr, "The right side of a register shift/scale must be an immediate");
+			return;
+		}
+
+		if (amount.mode == Addressing_Constant && amount.value.kind == ExactValue_Integer) {
+			i64 amt = exact_value_to_i64(amount.value);
+			if (be->op.kind == Token_Mul) {
+				if (amt <= 0 || (amt & (amt-1)) != 0) {
+					error(be->right, "A register scale using '*' must be a positive power of two, got %lld", cast(long long)amt);
+					return;
+				}
+			} else {
+				i64 max_shift = 63;
+				if (reg.type != nullptr && is_type_integer(reg.type)) {
+					max_shift = 8*cast(i64)type_size_of(reg.type) - 1;
+				}
+				if (amt < 0 || amt > max_shift) {
+					error(be->right, "A register shift amount must be within 0..=%lld, got %lld", cast(long long)max_shift, cast(long long)amt);
+					return;
+				}
+			}
+		}
+
+		operand->mode = reg.mode;
+		operand->type = reg.type;
+		return;
+	case_end;
 	}
 
 	{

src/llvm_backend_asm.cpp

@@ -1219,7 +1219,7 @@ struct lbAsmGenerate_arm64 : lbAsmGenerate {
 		}
 		auto slot = form.ops[i];
 		AsmOperandKind k = g_asm_arm64.kind_from_operand_type(slot);
-		if (k != AsmOperand_Register && k != AsmOperand_Register_Or_Memory) {
+		if (k != AsmOperand_Register && k != AsmOperand_Register_Or_Memory && k != AsmOperand_RegisterShift) {
 			return 0;
 		}
 		AsmRegClass cls = g_asm_arm64.operand_type_reg_class(slot);
@@ -1543,6 +1543,55 @@ struct lbAsmGenerate_arm64 : lbAsmGenerate {
 				break;
 			}
 		case_end;
+		case_ast_node(be, BinaryExpr, op);
+			// Shifted/scaled register operand -> `reg, <shift> #n`. `<<`=lsl, `>>`=lsr,
+			// and `*` is lsl by log2 of the (power-of-two) multiplier. Only produced on
+			// ARM64 (the checker rejects it elsewhere) and only for a *_SHIFTED slot.
+			char const *shift_name = nullptr;
+			switch (be->op.kind) {
+			case Token_Shl: shift_name = "lsl"; break;
+			case Token_Shr: shift_name = "lsr"; break;
+			case Token_Mul: shift_name = "lsl"; break;
+			default:
+				GB_PANIC("asm: unexpected register-shift operator '%.*s'", LIT(be->op.string));
+				break;
+			}
+
+			// The register takes the slot's own w/x modifier (arm64_slot_reg_modifier now
+			// covers RegisterShift), so recurse for it, then append the shift modifier.
+			this->write_operand(op_number, be->left, flags & ~WriteOperandFlag_PrintPrefixes);
+
+			Ast *amount = be->right;
+			if (amount->tav.mode == Addressing_Constant) {
+				i64 v = exact_value_to_i64(exact_value_to_integer(amount->tav.value));
+				i64 shift = v;
+				if (be->op.kind == Token_Mul) {
+					// reg * 2^k  ==  reg, lsl #k
+					shift = 0;
+					while (v > 1) {
+						v >>= 1;
+						shift++;
+					}
+				}
+				asm_string = gb_string_append_fmt(asm_string, ", %s #%lld", shift_name, cast(long long)shift);
+			} else {
+				// $-immediate shift amount. A runtime amount has no compile-time log2, so
+				// '*' must be a constant; '<<'/'>>' substitute the value via LLVM ($idx).
+				if (be->op.kind == Token_Mul) {
+					error(amount, "A '*' register scale needs a constant power-of-two amount");
+					break;
+				}
+				Entity *e  = entity_of_node(amount);
+				auto   *ed = entity_op(e);
+				if (ed == nullptr || ed->kind != AsmTemplateEntityDecl_Immediate) {
+					error(amount, "A register shift amount must be a constant or $-immediate");
+					break;
+				}
+				i32 idx = op_number[ed->total_index];
+				GB_ASSERT(idx >= 0);
+				asm_string = gb_string_append_fmt(asm_string, ", %s #$%d", shift_name, idx);
+			}
+		case_end;
 		default:
 			GB_PANIC("TODO(bill): write_operand for '%s'", expr_to_string(op));
 			break;

src/parser.cpp

@@ -2516,6 +2516,19 @@ gb_internal Ast *parse_asm_register(AstFile *f) {
 	return reg;
 }
 
+
+gb_internal bool asm_operand_is_shiftable(Ast *operand) {
+	if (operand == nullptr) {
+		return false;
+	}
+	switch (operand->kind) {
+	case Ast_AsmRegister:
+	case Ast_Ident:
+		return true;
+	}
+	return false;
+}
+
 gb_internal Ast *ast_asm_memory_term(AstFile *f, Token op, Ast *operand, Token scale_op, Ast *scale) {
 	Ast *term = alloc_ast_node(f, Ast_AsmMemoryTerm);
 	term->AsmMemoryTerm.op = op;
@@ -2712,7 +2725,10 @@ gb_internal Ast *parse_asm_operand(AstFile *f, bool allow_memory_operand) {
 	if (operand == nullptr) {
 		syntax_error(f->curr_token, "Invalid asm operand, found '%.*s'", LIT(f->curr_token.string));
 		advance_token(f);
-	} else if (f->curr_token.kind == Token_OpenBracket) {
+		return operand;
+	}
+
+	if (f->curr_token.kind == Token_OpenBracket) {
 		f->expr_level++;
 		Token open = expect_token(f, Token_OpenBracket);
 		Ast *index = parse_asm_operand(f, false);
@@ -2720,6 +2736,25 @@ gb_internal Ast *parse_asm_operand(AstFile *f, bool allow_memory_operand) {
 		f->expr_level--;
 		operand = ast_index_expr(f, operand, index, open, close);
 	}
+
+	if (allow_memory_operand && asm_operand_is_shiftable(operand)) {
+		switch (f->curr_token.kind) {
+		case Token_Mul:
+		case Token_Shl:
+		case Token_Shr:
+			{
+				Token op = advance_token(f);
+				Ast *amount = parse_asm_operand(f, false);
+				Ast *be = alloc_ast_node(f, Ast_BinaryExpr);
+				be->BinaryExpr.left = operand;
+				be->BinaryExpr.op = op;
+				be->BinaryExpr.right = amount;
+				operand = be;
+			}
+			break;
+		}
+	}
+
 	return operand;
 }