Implement instrumentation pass

src/check_decl.cpp

@@ -910,24 +910,24 @@ gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d) {
 	e->Procedure.entry_point_only = ac.entry_point_only;
 	e->Procedure.is_export = ac.is_export;
 
-	bool no_instrumentation = false;
+	bool has_instrumentation = false;
 	if (pl->body == nullptr) {
-		no_instrumentation = true;
+		has_instrumentation = false;
 		if (ac.no_instrumentation != Instrumentation_Default) {
 			error(e->token, "@(no_instrumentation) is not allowed on foreign procedures");
 		}
 	} else {
-		if (e->file) {
-			no_instrumentation = (e->file->flags & AstFile_NoInstrumentation) != 0;
+		AstFile *file = e->token.pos.file_id ? global_files[e->token.pos.file_id] : nullptr;
+		if (file) {
+			has_instrumentation = (file->flags & AstFile_NoInstrumentation) == 0;
 		}
 
 		switch (ac.no_instrumentation) {
-		case Instrumentation_Enabled:  no_instrumentation = false; break;
+		case Instrumentation_Enabled:  has_instrumentation = true; break;
 		case Instrumentation_Default:  break;
-		case Instrumentation_Disabled: no_instrumentation = true;  break;
+		case Instrumentation_Disabled: has_instrumentation = false;  break;
 		}
 	}
-	e->Procedure.no_instrumentation = no_instrumentation;
 
 	auto const is_valid_instrumentation_call = [](Type *type) -> bool {
 		if (type == nullptr || type->kind != Type_Proc) {
@@ -949,6 +949,9 @@ gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d) {
 
 	if (ac.instrumentation_enter && ac.instrumentation_exit) {
 		error(e->token, "A procedure cannot be marked with both @(instrumentation_enter) and @(instrumentation_exit)");
+
+		has_instrumentation = false;
+		e->flags |= EntityFlag_Require;
 	} else if (ac.instrumentation_enter) {
 		if (!is_valid_instrumentation_call(e->type)) {
 			gbString s = type_to_string(e->type);
@@ -961,6 +964,9 @@ gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d) {
 		} else {
 			ctx->info->instrumentation_enter_entity = e;
 		}
+
+		has_instrumentation = false;
+		e->flags |= EntityFlag_Require;
 	} else if (ac.instrumentation_exit) {
 		if (!is_valid_instrumentation_call(e->type)) {
 			gbString s = type_to_string(e->type);
@@ -973,8 +979,14 @@ gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d) {
 		} else {
 			ctx->info->instrumentation_exit_entity = e;
 		}
+
+		has_instrumentation = false;
+		e->flags |= EntityFlag_Require;
 	}
 
+	e->Procedure.has_instrumentation = has_instrumentation;
+
+
 	e->deprecated_message = ac.deprecated_message;
 	e->warning_message = ac.warning_message;
 	ac.link_name = handle_link_name(ctx, e->token, ac.link_name, ac.link_prefix);

src/entity.cpp

@@ -251,7 +251,7 @@ struct Entity {
 			bool    generated_from_polymorphic : 1;
 			bool    target_feature_disabled    : 1;
 			bool    entry_point_only           : 1;
-			bool    no_instrumentation         : 1;
+			bool    has_instrumentation        : 1;
 			String  target_feature;
 		} Procedure;
 		struct {

src/llvm_backend.cpp

@@ -1497,8 +1497,6 @@ gb_internal WORKER_TASK_PROC(lb_llvm_module_pass_worker_proc) {
 	auto passes = array_make<char const *>(heap_allocator(), 0, 64);
 	defer (array_free(&passes));
 
-
-
 	LLVMPassBuilderOptionsRef pb_options = LLVMCreatePassBuilderOptions();
 	defer (LLVMDisposePassBuilderOptions(pb_options));
 

src/llvm_backend.hpp

@@ -563,7 +563,7 @@ gb_internal LLVMTypeRef OdinLLVMGetVectorElementType(LLVMTypeRef type);
 
 gb_internal String lb_filepath_ll_for_module(lbModule *m);
 
-
+gb_internal LLVMTypeRef lb_type_internal_for_procedures_raw(lbModule *m, Type *type);
 
 gb_internal LLVMTypeRef llvm_array_type(LLVMTypeRef ElementType, uint64_t ElementCount) {
 #if LB_USE_NEW_PASS_SYSTEM

src/llvm_backend_general.cpp

@@ -2348,6 +2348,15 @@ gb_internal LLVMAttributeRef lb_create_enum_attribute(LLVMContextRef ctx, char c
 	return LLVMCreateEnumAttribute(ctx, kind, value);
 }
 
+gb_internal LLVMAttributeRef lb_create_string_attribute(LLVMContextRef ctx, String const &key, String const &value) {
+	LLVMAttributeRef attr = LLVMCreateStringAttribute(
+		ctx,
+		cast(char const *)key.text,   cast(unsigned)key.len,
+		cast(char const *)value.text, cast(unsigned)value.len);
+	return attr;
+}
+
+
 gb_internal void lb_add_proc_attribute_at_index(lbProcedure *p, isize index, char const *name, u64 value) {
 	LLVMAttributeRef attr = lb_create_enum_attribute(p->module->ctx, name, value);
 	GB_ASSERT(attr != nullptr);
@@ -2361,6 +2370,10 @@ gb_internal void lb_add_proc_attribute_at_index(lbProcedure *p, isize index, cha
 gb_internal void lb_add_attribute_to_proc(lbModule *m, LLVMValueRef proc_value, char const *name, u64 value=0) {
 	LLVMAddAttributeAtIndex(proc_value, LLVMAttributeIndex_FunctionIndex, lb_create_enum_attribute(m->ctx, name, value));
 }
+gb_internal void lb_add_attribute_to_proc_with_string(lbModule *m, LLVMValueRef proc_value, String const &name, String const &value) {
+	LLVMAttributeRef attr = lb_create_string_attribute(m->ctx, name, value);
+	LLVMAddAttributeAtIndex(proc_value, LLVMAttributeIndex_FunctionIndex, attr);
+}
 
 
 

src/llvm_backend_opt.cpp

@@ -380,6 +380,80 @@ gb_internal void lb_run_remove_dead_instruction_pass(lbProcedure *p) {
 	}
 }
 
+gb_internal LLVMValueRef lb_run_instrumentation_pass_insert_call(lbProcedure *p, Entity *entity, LLVMBuilderRef dummy_builder) {
+	lbModule *m = p->module;
+
+	lbValue cc = lb_find_procedure_value_from_entity(m, entity);
+
+	LLVMValueRef args[2] = {};
+	args[0] = p->value;
+
+	LLVMValueRef returnaddress_args[1] = {};
+
+	returnaddress_args[0] = LLVMConstInt(LLVMInt32TypeInContext(m->ctx), 0, false);
+
+	char const *instrinsic_name = "llvm.returnaddress";
+	unsigned id = LLVMLookupIntrinsicID(instrinsic_name, gb_strlen(instrinsic_name));
+	GB_ASSERT_MSG(id != 0, "Unable to find %s", instrinsic_name);
+	LLVMValueRef ip = LLVMGetIntrinsicDeclaration(m->mod, id, nullptr, 0);
+	LLVMTypeRef call_type = LLVMIntrinsicGetType(m->ctx, id, nullptr, 0);
+	args[1] = LLVMBuildCall2(dummy_builder, call_type, ip, returnaddress_args, gb_count_of(returnaddress_args), "");
+
+	LLVMTypeRef fnp = lb_type_internal_for_procedures_raw(p->module, entity->type);
+	return LLVMBuildCall2(dummy_builder, fnp, cc.value, args, 2, "");
+}
+
+
+gb_internal void lb_run_instrumentation_pass(lbProcedure *p) {
+	lbModule *m = p->module;
+	Entity *enter = m->info->instrumentation_enter_entity;
+	Entity *exit  = m->info->instrumentation_exit_entity;
+	if (enter == nullptr || exit == nullptr) {
+		return;
+	}
+	if (!(p->entity &&
+	      p->entity->kind == Entity_Procedure &&
+	      p->entity->Procedure.has_instrumentation)) {
+		return;
+	}
+
+#define LLVM_V_NAME(x) x, cast(unsigned)(gb_count_of(x)-1)
+
+	LLVMBuilderRef dummy_builder = LLVMCreateBuilderInContext(m->ctx);
+	defer (LLVMDisposeBuilder(dummy_builder));
+
+	LLVMBasicBlockRef entry_bb = p->entry_block->block;
+	LLVMPositionBuilder(dummy_builder, entry_bb, LLVMGetFirstInstruction(entry_bb));
+	lb_run_instrumentation_pass_insert_call(p, enter, dummy_builder);
+	LLVMRemoveStringAttributeAtIndex(p->value, LLVMAttributeIndex_FunctionIndex, LLVM_V_NAME("instrument-function-entry"));
+
+	unsigned bb_count = LLVMCountBasicBlocks(p->value);
+	LLVMBasicBlockRef *bbs = gb_alloc_array(temporary_allocator(), LLVMBasicBlockRef, bb_count);
+	LLVMGetBasicBlocks(p->value, bbs);
+	for (unsigned i = 0; i < bb_count; i++) {
+		LLVMBasicBlockRef bb = bbs[i];
+		LLVMValueRef terminator = LLVMGetBasicBlockTerminator(bb);
+		if (terminator == nullptr ||
+		    !LLVMIsAReturnInst(terminator)) {
+			continue;
+		}
+
+		// TODO(bill): getTerminatingMustTailCall()
+		// If T is preceded by a musttail call, that's the real terminator.
+		// if (CallInst *CI = BB.getTerminatingMustTailCall())
+		// 	T = CI;
+
+
+		LLVMPositionBuilderBefore(dummy_builder, terminator);
+		lb_run_instrumentation_pass_insert_call(p, exit, dummy_builder);
+	}
+
+	LLVMRemoveStringAttributeAtIndex(p->value, LLVMAttributeIndex_FunctionIndex, LLVM_V_NAME("instrument-function-exit"));
+
+#undef LLVM_V_NAME
+}
+
+
 
 gb_internal void lb_run_function_pass_manager(LLVMPassManagerRef fpm, lbProcedure *p, lbFunctionPassManagerKind pass_manager_kind) {
 	if (p == nullptr) {
@@ -401,6 +475,7 @@ gb_internal void lb_run_function_pass_manager(LLVMPassManagerRef fpm, lbProcedur
 	    }
 	    break;
 	}
+	lb_run_instrumentation_pass(p);
 
 	LLVMRunFunctionPassManager(fpm, p->value);
 }
@@ -552,3 +627,5 @@ gb_internal void lb_run_remove_unused_globals_pass(lbModule *m) {
 		}
 	}
 }
+
+

src/llvm_backend_proc.cpp

@@ -329,6 +329,18 @@ gb_internal lbProcedure *lb_create_procedure(lbModule *m, Entity *entity, bool i
 		}
 	}
 
+	if (p->body && entity->Procedure.has_instrumentation) {
+		Entity *instrumentation_enter = m->info->instrumentation_enter_entity;
+		Entity *instrumentation_exit  = m->info->instrumentation_exit_entity;
+		if (instrumentation_enter && instrumentation_exit) {
+			String enter = lb_get_entity_name(m, instrumentation_enter);
+			String exit  = lb_get_entity_name(m, instrumentation_exit);
+
+			lb_add_attribute_to_proc_with_string(m, p->value, make_string_c("instrument-function-entry"), enter);
+			lb_add_attribute_to_proc_with_string(m, p->value, make_string_c("instrument-function-exit"),  exit);
+		}
+	}
+
 	lbValue proc_value = {p->value, p->type};
 	lb_add_entity(m, entity,  proc_value);
 	lb_add_member(m, p->name, proc_value);