Merge pull request #2055 from odin-lang/map-index-internal

Map Internals Improvements

5 files changed, 88 insertions, 84 deletions

diff --git a/src/llvm_backend.cpp b/src/llvm_backend.cpp
index aa901d22f..142ecc348 100644
--- a/src/llvm_backend.cpp
+++ b/src/llvm_backend.cpp
@@ -502,48 +502,58 @@ lbValue lb_generate_anonymous_proc_lit(lbModule *m, String const &prefix_name, A
 
 lbValue lb_gen_map_header(lbProcedure *p, lbValue map_val_ptr, Type *map_type) {
 	GB_ASSERT_MSG(is_type_pointer(map_val_ptr.type), "%s", type_to_string(map_val_ptr.type));
-	lbAddr h = lb_add_local_generated(p, t_map_header, false); // all the values will be initialzed later
 	map_type = base_type(map_type);
 	GB_ASSERT(map_type->kind == Type_Map);
 
-	Type *key_type = map_type->Map.key;
-	Type *val_type = map_type->Map.value;
-	gb_unused(val_type);
+	lbAddr h = {};
+	lbAddr *found = map_get(&p->map_header_cache, map_val_ptr.value);
+	if (found != nullptr) {
+		h = *found;
+	} else {
+		h = lb_add_local_generated(p, t_map_header, false); // all the values will be initialzed later
 
-	GB_ASSERT(map_type->Map.entry_type->kind == Type_Struct);
-	map_type->Map.entry_type->cached_size = -1;
-	map_type->Map.entry_type->Struct.are_offsets_set = false;
-	
-	i64 entry_size   = type_size_of  (map_type->Map.entry_type);
-	i64 entry_align  = type_align_of (map_type->Map.entry_type);
-	
-	i64 key_offset = type_offset_of(map_type->Map.entry_type, 2);
-	i64 key_size   = type_size_of  (map_type->Map.key);
+		Type *key_type = map_type->Map.key;
+		Type *val_type = map_type->Map.value;
+		gb_unused(val_type);
+
+		GB_ASSERT(map_type->Map.entry_type->kind == Type_Struct);
+		map_type->Map.entry_type->cached_size = -1;
+		map_type->Map.entry_type->Struct.are_offsets_set = false;
+
+		i64 entry_size   = type_size_of  (map_type->Map.entry_type);
+		i64 entry_align  = type_align_of (map_type->Map.entry_type);
+
+		i64 key_offset = type_offset_of(map_type->Map.entry_type, 2);
+		i64 key_size   = type_size_of  (map_type->Map.key);
+
+		i64 value_offset = type_offset_of(map_type->Map.entry_type, 3);
+		i64 value_size   = type_size_of  (map_type->Map.value);
 
-	i64 value_offset = type_offset_of(map_type->Map.entry_type, 3);
-	i64 value_size   = type_size_of  (map_type->Map.value);
-	
-	
-	Type *map_header_base = base_type(t_map_header);
-	GB_ASSERT(map_header_base->Struct.fields.count == 8);
-	Type *raw_map_ptr_type = map_header_base->Struct.fields[0]->type;
-	LLVMValueRef const_values[8] = {};
-	const_values[0] = LLVMConstNull(lb_type(p->module, raw_map_ptr_type));
-	const_values[1] = lb_get_equal_proc_for_type(p->module, key_type)    .value;
-	const_values[2] = lb_const_int(p->module, t_int,        entry_size)  .value;
-	const_values[3] = lb_const_int(p->module, t_int,        entry_align) .value;
-	const_values[4] = lb_const_int(p->module, t_uintptr,    key_offset)  .value;
-	const_values[5] = lb_const_int(p->module, t_int,        key_size)    .value;
-	const_values[6] = lb_const_int(p->module, t_uintptr,    value_offset).value;
-	const_values[7] = lb_const_int(p->module, t_int,        value_size)  .value;
-	
-	LLVMValueRef const_value = llvm_const_named_struct(p->module, t_map_header, const_values, gb_count_of(const_values));
-	LLVMBuildStore(p->builder, const_value, h.addr.value);
-	
-	// NOTE(bill): Removes unnecessary allocation if split gep
-	lbValue gep0 = lb_emit_struct_ep(p, h.addr, 0);
-	lbValue m = lb_emit_conv(p, map_val_ptr, type_deref(gep0.type));
-	lb_emit_store(p, gep0, m);
+
+		Type *map_header_base = base_type(t_map_header);
+		GB_ASSERT(map_header_base->Struct.fields.count == 8);
+		Type *raw_map_ptr_type = map_header_base->Struct.fields[0]->type;
+		LLVMValueRef const_values[8] = {};
+		const_values[0] = LLVMConstNull(lb_type(p->module, raw_map_ptr_type));
+		const_values[1] = lb_get_equal_proc_for_type(p->module, key_type)    .value;
+		const_values[2] = lb_const_int(p->module, t_int,        entry_size)  .value;
+		const_values[3] = lb_const_int(p->module, t_int,        entry_align) .value;
+		const_values[4] = lb_const_int(p->module, t_uintptr,    key_offset)  .value;
+		const_values[5] = lb_const_int(p->module, t_int,        key_size)    .value;
+		const_values[6] = lb_const_int(p->module, t_uintptr,    value_offset).value;
+		const_values[7] = lb_const_int(p->module, t_int,        value_size)  .value;
+
+		LLVMValueRef const_value = llvm_const_named_struct(p->module, t_map_header, const_values, gb_count_of(const_values));
+		LLVMBuildStore(p->builder, const_value, h.addr.value);
+
+		// NOTE(bill): Removes unnecessary allocation if split gep
+		lbValue gep0 = lb_emit_struct_ep(p, h.addr, 0);
+		lbValue m = lb_emit_conv(p, map_val_ptr, type_deref(gep0.type));
+		lb_emit_store(p, gep0, m);
+
+
+		map_set(&p->map_header_cache, map_val_ptr.value, h);
+	}
 
 	return lb_addr_load(p, h);
 }
@@ -595,14 +605,12 @@ lbValue lb_const_hash(lbModule *m, lbValue key, Type *key_type) {
 	return hashed_key;
 }
 
-lbValue lb_gen_map_hash(lbProcedure *p, lbValue key, Type *key_type) {
-	lbAddr v = lb_add_local_generated(p, t_map_hash, true);
-	lbValue vp = lb_addr_get_ptr(p, v);
-	key = lb_emit_conv(p, key, key_type);
-
+lbValue lb_gen_map_key_hash(lbProcedure *p, lbValue key, Type *key_type, lbValue *key_ptr_) {
 	lbValue key_ptr = lb_address_from_load_or_generate_local(p, key);
 	key_ptr = lb_emit_conv(p, key_ptr, t_rawptr);
 
+	if (key_ptr_) *key_ptr_ = key_ptr;
+
 	lbValue hashed_key = lb_const_hash(p->module, key, key_type);
 	if (hashed_key.value == nullptr) {
 		lbValue hasher = lb_get_hasher_proc_for_type(p->module, key_type);
@@ -613,10 +621,7 @@ lbValue lb_gen_map_hash(lbProcedure *p, lbValue key, Type *key_type) {
 		hashed_key = lb_emit_call(p, hasher, args);
 	}
 
-	lb_emit_store(p, lb_emit_struct_ep(p, vp, 0), hashed_key);
-	lb_emit_store(p, lb_emit_struct_ep(p, vp, 1), key_ptr);
-
-	return lb_addr_load(p, v);
+	return hashed_key;
 }
 
 void lb_insert_dynamic_map_key_and_value(lbProcedure *p, lbAddr addr, Type *map_type,
@@ -625,17 +630,19 @@ void lb_insert_dynamic_map_key_and_value(lbProcedure *p, lbAddr addr, Type *map_
 	GB_ASSERT(map_type->kind == Type_Map);
 
 	lbValue h = lb_gen_map_header(p, addr.addr, map_type);
-	lbValue key = lb_gen_map_hash(p, map_key, map_type->Map.key);
+	lbValue key_ptr = {};
+	lbValue key_hash = lb_gen_map_key_hash(p, map_key, map_type->Map.key, &key_ptr);
 	lbValue v = lb_emit_conv(p, map_value, map_type->Map.value);
 
 	lbAddr value_addr = lb_add_local_generated(p, v.type, false);
 	lb_addr_store(p, value_addr, v);
 
-	auto args = array_make<lbValue>(permanent_allocator(), 4);
+	auto args = array_make<lbValue>(permanent_allocator(), 5);
 	args[0] = h;
-	args[1] = key;
-	args[2] = lb_emit_conv(p, value_addr.addr, t_rawptr);
-	args[3] = lb_emit_source_code_location(p, node);
+	args[1] = key_hash;
+	args[2] = key_ptr;
+	args[3] = lb_emit_conv(p, value_addr.addr, t_rawptr);
+	args[4] = lb_emit_source_code_location(p, node);
 	lb_emit_runtime_call(p, "__dynamic_map_set", args);
 }
 
diff --git a/src/llvm_backend.hpp b/src/llvm_backend.hpp
index 79f0f37e7..d622f3661 100644
--- a/src/llvm_backend.hpp
+++ b/src/llvm_backend.hpp
@@ -308,6 +308,7 @@ struct lbProcedure {
 
 	PtrMap<Ast *, lbValue> selector_values;
 	PtrMap<Ast *, lbAddr>  selector_addr;
+	PtrMap<LLVMValueRef, lbAddr> map_header_cache;
 };
 
 
@@ -444,7 +445,7 @@ String lb_get_const_string(lbModule *m, lbValue value);
 lbValue lb_generate_local_array(lbProcedure *p, Type *elem_type, i64 count, bool zero_init=true);
 lbValue lb_generate_global_array(lbModule *m, Type *elem_type, i64 count, String prefix, i64 id);
 lbValue lb_gen_map_header(lbProcedure *p, lbValue map_val_ptr, Type *map_type);
-lbValue lb_gen_map_hash(lbProcedure *p, lbValue key, Type *key_type);
+lbValue lb_gen_map_key_hash(lbProcedure *p, lbValue key, Type *key_type, lbValue *key_ptr_);
 void    lb_insert_dynamic_map_key_and_value(lbProcedure *p, lbAddr addr, Type *map_type, lbValue map_key, lbValue map_value, Ast *node);
 
 lbValue lb_find_procedure_value_from_entity(lbModule *m, Entity *e);
diff --git a/src/llvm_backend_expr.cpp b/src/llvm_backend_expr.cpp
index 7c92c517c..3ab73a27b 100644
--- a/src/llvm_backend_expr.cpp
+++ b/src/llvm_backend_expr.cpp
@@ -1423,15 +1423,9 @@ lbValue lb_build_binary_expr(lbProcedure *p, Ast *expr) {
 			switch (rt->kind) {
 			case Type_Map:
 				{
-					lbValue addr = lb_address_from_load_or_generate_local(p, right);
-					lbValue h = lb_gen_map_header(p, addr, rt);
-					lbValue key = lb_gen_map_hash(p, left, rt->Map.key);
-
-					auto args = array_make<lbValue>(permanent_allocator(), 2);
-					args[0] = h;
-					args[1] = key;
-
-					lbValue ptr = lb_emit_runtime_call(p, "__dynamic_map_get", args);
+					lbValue map_ptr = lb_address_from_load_or_generate_local(p, right);
+					lbValue key = left;
+					lbValue ptr = lb_internal_dynamic_map_get_ptr(p, map_ptr, key);
 					if (be->op.kind == Token_in) {
 						return lb_emit_conv(p, lb_emit_comp_against_nil(p, Token_NotEq, ptr), t_bool);
 					} else {
diff --git a/src/llvm_backend_general.cpp b/src/llvm_backend_general.cpp
index 1f8fccdcb..55b09cbfc 100644
--- a/src/llvm_backend_general.cpp
+++ b/src/llvm_backend_general.cpp
@@ -383,6 +383,21 @@ Type *lb_addr_type(lbAddr const &addr) {
 	return type_deref(addr.addr.type);
 }
 
+lbValue lb_internal_dynamic_map_get_ptr(lbProcedure *p, lbValue const &map_ptr, lbValue const &key) {
+	Type *map_type = base_type(type_deref(map_ptr.type));
+	lbValue h = lb_gen_map_header(p, map_ptr, map_type);
+
+	lbValue key_ptr = {};
+	auto args = array_make<lbValue>(permanent_allocator(), 3);
+	args[0] = h;
+	args[1] = lb_gen_map_key_hash(p, key, map_type->Map.key, &key_ptr);
+	args[2] = key_ptr;
+
+	lbValue ptr = lb_emit_runtime_call(p, "__dynamic_map_get", args);
+
+	return lb_emit_conv(p, ptr, alloc_type_pointer(map_type->Map.value));
+}
+
 lbValue lb_addr_get_ptr(lbProcedure *p, lbAddr const &addr) {
 	if (addr.addr.value == nullptr) {
 		GB_PANIC("Illegal addr -> nullptr");
@@ -390,19 +405,8 @@ lbValue lb_addr_get_ptr(lbProcedure *p, lbAddr const &addr) {
 	}
 
 	switch (addr.kind) {
-	case lbAddr_Map: {
-		Type *map_type = base_type(addr.map.type);
-		lbValue h = lb_gen_map_header(p, addr.addr, map_type);
-		lbValue key = lb_gen_map_hash(p, addr.map.key, map_type->Map.key);
-
-		auto args = array_make<lbValue>(permanent_allocator(), 2);
-		args[0] = h;
-		args[1] = key;
-
-		lbValue ptr = lb_emit_runtime_call(p, "__dynamic_map_get", args);
-
-		return lb_emit_conv(p, ptr, alloc_type_pointer(map_type->Map.value));
-	}
+	case lbAddr_Map:
+		return lb_internal_dynamic_map_get_ptr(p, addr.addr, addr.map.key);
 
 	case lbAddr_RelativePointer: {
 		Type *rel_ptr = base_type(lb_addr_type(addr));
@@ -1059,16 +1063,11 @@ lbValue lb_addr_load(lbProcedure *p, lbAddr const &addr) {
 
 
 	} else if (addr.kind == lbAddr_Map) {
-		Type *map_type = base_type(addr.map.type);
+		Type *map_type = base_type(type_deref(addr.addr.type));
+		GB_ASSERT(map_type->kind == Type_Map);
 		lbAddr v = lb_add_local_generated(p, map_type->Map.lookup_result_type, true);
-		lbValue h = lb_gen_map_header(p, addr.addr, map_type);
-		lbValue key = lb_gen_map_hash(p, addr.map.key, map_type->Map.key);
-
-		auto args = array_make<lbValue>(permanent_allocator(), 2);
-		args[0] = h;
-		args[1] = key;
 
-		lbValue ptr = lb_emit_runtime_call(p, "__dynamic_map_get", args);
+		lbValue ptr = lb_internal_dynamic_map_get_ptr(p, addr.addr, addr.map.key);
 		lbValue ok = lb_emit_conv(p, lb_emit_comp_against_nil(p, Token_NotEq, ptr), t_bool);
 		lb_emit_store(p, lb_emit_struct_ep(p, v.addr, 1), ok);
 
diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
index 8bbbb0c56..17ed9c2a6 100644
--- a/src/llvm_backend_proc.cpp
+++ b/src/llvm_backend_proc.cpp
@@ -121,8 +121,9 @@ lbProcedure *lb_create_procedure(lbModule *m, Entity *entity, bool ignore_body)
 	p->branch_blocks.allocator = a;
 	p->context_stack.allocator = a;
 	p->scope_stack.allocator   = a;
-	map_init(&p->selector_values, a, 0);
-	map_init(&p->selector_addr,   a, 0);
+	map_init(&p->selector_values,  a, 0);
+	map_init(&p->selector_addr,    a, 0);
+	map_init(&p->map_header_cache, a, 0);
 
 	if (p->is_foreign) {
 		lb_add_foreign_library_path(p->module, entity->Procedure.foreign_library);
@@ -380,6 +381,8 @@ lbProcedure *lb_create_dummy_procedure(lbModule *m, String link_name, Type *type
 		lb_add_proc_attribute_at_index(p, offset+parameter_index, "nocapture");
 	}
 
+	map_init(&p->map_header_cache, heap_allocator(), 0);
+
 	return p;
 }