checker: support integer and voidptr key types for maps (#7503)

vlib/builtin/map.v

@@ -139,18 +139,10 @@ fn (d &DenseArray) has_index(i int) bool {
 	return d.deletes == 0 || unsafe { d.all_deleted[i] } == 0
 }
 
-[inline]
-fn (d &DenseArray) clone_key(dest voidptr, pkey voidptr) {
-	unsafe {
-		s := (*&string(pkey)).clone()
-		C.memcpy(dest, &s, d.key_bytes)
-	}
-}
-
-// Push element to array and return index
+// Make space to append an element and return index
 // The growth-factor is roughly 1.125 `(x + (x >> 3))`
 [inline]
-fn (mut d DenseArray) push(key voidptr, value voidptr) int {
+fn (mut d DenseArray) expand() int {
 	if d.cap == d.len {
 		d.cap += d.cap >> 3
 		unsafe {
@@ -166,9 +158,6 @@ fn (mut d DenseArray) push(key voidptr, value voidptr) int {
 		if d.deletes != 0 {
 			d.all_deleted[push_index] = 0
 		}
-		ptr := d.key(push_index)
-		d.clone_key(ptr, key)
-		C.memcpy(byteptr(ptr) + d.key_bytes, value, d.value_bytes)
 	}
 	d.len++
 	return push_index
@@ -201,6 +190,14 @@ fn (mut d DenseArray) zeros_to_end() {
 	}
 }
 
+type MapHashFn = fn (voidptr) u64
+
+type MapEqFn = fn (voidptr, voidptr) bool
+
+type MapCloneFn = fn (voidptr, voidptr)
+
+type MapFreeFn = fn (voidptr)
+
 pub struct map {
 	// Number of bytes of a key
 	key_bytes       int
@@ -222,11 +219,95 @@ mut:
 	// Extra metas that allows for no ranging when incrementing
 	// index in the hashmap
 	extra_metas     u32
+	has_string_keys bool
+	hash_fn         MapHashFn
+	key_eq_fn       MapEqFn
+	clone_fn        MapCloneFn
+	free_fn         MapFreeFn
 pub mut:
 	// Number of key-values currently in the hashmap
 	len             int
 }
 
+fn map_hash_string(pkey voidptr) u64 {
+	key := *&string(pkey)
+	return hash.wyhash_c(key.str, u64(key.len), 0)
+}
+
+fn map_hash_int_1(pkey voidptr) u64 {
+	return hash.wyhash64_c(*&byte(pkey), 0)
+}
+
+fn map_hash_int_2(pkey voidptr) u64 {
+	return hash.wyhash64_c(*&u16(pkey), 0)
+}
+
+fn map_hash_int_4(pkey voidptr) u64 {
+	return hash.wyhash64_c(*&u32(pkey), 0)
+}
+
+fn map_hash_int_8(pkey voidptr) u64 {
+	return hash.wyhash64_c(*&u64(pkey), 0)
+}
+
+fn map_eq_string(a voidptr, b voidptr) bool {
+	return fast_string_eq(*&string(a), *&string(b))
+}
+
+fn map_eq_int_1(a voidptr, b voidptr) bool {
+	return *&byte(a) == *&byte(b)
+}
+
+fn map_eq_int_2(a voidptr, b voidptr) bool {
+	return *&u16(a) == *&u16(b)
+}
+
+fn map_eq_int_4(a voidptr, b voidptr) bool {
+	return *&u32(a) == *&u32(b)
+}
+
+fn map_eq_int_8(a voidptr, b voidptr) bool {
+	return *&u64(a) == *&u64(b)
+}
+
+fn map_clone_string(dest voidptr, pkey voidptr) {
+	unsafe {
+		s := *&string(pkey)
+		(*&string(dest)) = s.clone()
+	}
+}
+
+fn map_clone_int_1(dest voidptr, pkey voidptr) {
+	unsafe {
+		*&byte(dest) = *&byte(pkey)
+	}
+}
+
+fn map_clone_int_2(dest voidptr, pkey voidptr) {
+	unsafe {
+		*&u16(dest) = *&u16(pkey)
+	}
+}
+
+fn map_clone_int_4(dest voidptr, pkey voidptr) {
+	unsafe {
+		*&u32(dest) = *&u32(pkey)
+	}
+}
+
+fn map_clone_int_8(dest voidptr, pkey voidptr) {
+	unsafe {
+		*&u64(dest) = *&u64(pkey)
+	}
+}
+
+fn map_free_string(pkey voidptr) {
+	(*&string(pkey)).free()
+}
+
+fn map_free_nop(_ voidptr) {
+}
+
 // bootstrap
 fn new_map_1(value_bytes int) map {
 	return new_map(int(sizeof(string)), value_bytes)
@@ -234,6 +315,45 @@ fn new_map_1(value_bytes int) map {
 
 fn new_map(key_bytes int, value_bytes int) map {
 	metasize := int(sizeof(u32) * (init_capicity + extra_metas_inc))
+	// for now assume anything bigger than a pointer is a string
+	has_string_keys := key_bytes > sizeof(voidptr)
+	mut hash_fn := MapHashFn(0)
+	mut key_eq_fn := MapEqFn(0)
+	mut clone_fn := MapCloneFn(0)
+	match key_bytes {
+		// assume non-string keys are bitwise comparable
+		1 {
+			hash_fn = &map_hash_int_1
+			key_eq_fn = &map_eq_int_1
+			clone_fn = &map_clone_int_1
+		}
+		2 {
+			hash_fn = &map_hash_int_2
+			key_eq_fn = &map_eq_int_2
+			clone_fn = &map_clone_int_2
+		}
+		4 {
+			hash_fn = &map_hash_int_4
+			key_eq_fn = &map_eq_int_4
+			clone_fn = &map_clone_int_4
+		}
+		8 {
+			hash_fn = &map_hash_int_8
+			key_eq_fn = &map_eq_int_8
+			clone_fn = &map_clone_int_8
+		}
+		else {
+			hash_fn = &map_hash_string
+			key_eq_fn = &map_eq_string
+			clone_fn = &map_clone_string
+		}
+	}
+	mut free_fn := MapFreeFn(0)
+	if has_string_keys {
+		free_fn = &map_free_string
+	} else {
+		free_fn = &map_free_nop
+	}
 	return map{
 		key_bytes: key_bytes
 		value_bytes: value_bytes
@@ -244,6 +364,11 @@ fn new_map(key_bytes int, value_bytes int) map {
 		metas: &u32(vcalloc(metasize))
 		extra_metas: extra_metas_inc
 		len: 0
+		has_string_keys: has_string_keys
+		hash_fn: hash_fn
+		key_eq_fn: key_eq_fn
+		clone_fn: clone_fn
+		free_fn: free_fn
 	}
 }
 
@@ -266,25 +391,14 @@ fn new_map_init_1(n int, key_bytes int, value_bytes int, keys voidptr, values vo
 	return out
 }
 
-[inline]
-fn (m &map) keys_eq(a voidptr, b voidptr) bool {
-	// assume string for now
-	return fast_string_eq(*&string(a), *&string(b))
-}
-
 [inline]
 fn (m &map) key_to_index(pkey voidptr) (u32, u32) {
-	key := *&string(pkey)
-	hash := hash.wyhash_c(key.str, u64(key.len), 0)
+	hash := m.hash_fn(pkey)
 	index := hash & m.even_index
 	meta := ((hash >> m.shift) & hash_mask) | probe_inc
 	return u32(index), u32(meta)
 }
 
-fn (m &map) free_key(pkey voidptr) {
-	(*&string(pkey)).free()
-}
-
 [inline]
 fn (m &map) meta_less(_index u32, _metas u32) (u32, u32) {
 	mut index := _index
@@ -359,7 +473,7 @@ fn (mut m map) set_1(key voidptr, value voidptr) {
 	for meta == unsafe { m.metas[index] } {
 		kv_index := int(unsafe { m.metas[index + 1] })
 		pkey := unsafe { m.key_values.key(kv_index) }
-		if m.keys_eq(key, pkey) {
+		if m.key_eq_fn(key, pkey) {
 			unsafe {
 				pval := byteptr(pkey) + m.key_bytes
 				C.memcpy(pval, value, m.value_bytes)
@@ -369,7 +483,12 @@ fn (mut m map) set_1(key voidptr, value voidptr) {
 		index += 2
 		meta += probe_inc
 	}
-	kv_index := m.key_values.push(key, value)
+	kv_index := m.key_values.expand()
+	unsafe {
+		pkey := m.key_values.key(kv_index)
+		m.clone_fn(pkey, key)
+		C.memcpy(byteptr(pkey) + m.key_bytes, value, m.value_bytes)
+	}
 	m.meta_greater(index, meta, u32(kv_index))
 	m.len++
 }
@@ -449,7 +568,7 @@ fn (mut m map) get_and_set_1(key voidptr, zero voidptr) voidptr {
 			if meta == unsafe { m.metas[index] } {
 				kv_index := int(unsafe { m.metas[index + 1] })
 				pkey := unsafe { m.key_values.key(kv_index) }
-				if m.keys_eq(key, pkey) {
+				if m.key_eq_fn(key, pkey) {
 					return unsafe { byteptr(pkey) + m.key_values.key_bytes }
 				}
 			}
@@ -479,7 +598,7 @@ fn (m &map) get_1(key voidptr, zero voidptr) voidptr {
 		if meta == unsafe { m.metas[index] } {
 			kv_index := int(unsafe { m.metas[index + 1] })
 			pkey := unsafe { m.key_values.key(kv_index) }
-			if m.keys_eq(key, pkey) {
+			if m.key_eq_fn(key, pkey) {
 				return unsafe { byteptr(pkey) + m.key_values.key_bytes }
 			}
 		}
@@ -503,7 +622,7 @@ fn (m &map) exists_1(key voidptr) bool {
 		if meta == unsafe { m.metas[index] } {
 			kv_index := int(unsafe { m.metas[index + 1] })
 			pkey := unsafe { m.key_values.key(kv_index) }
-			if m.keys_eq(key, pkey) {
+			if m.key_eq_fn(key, pkey) {
 				return true
 			}
 		}
@@ -539,7 +658,7 @@ pub fn (mut m map) delete_1(key voidptr) {
 	for meta == unsafe { m.metas[index] } {
 		kv_index := int(unsafe { m.metas[index + 1] })
 		pkey := unsafe { m.key_values.key(kv_index) }
-		if m.keys_eq(key, pkey) {
+		if m.key_eq_fn(key, pkey) {
 			for (unsafe { m.metas[index + 2] } >> hashbits) > 1 {
 				unsafe {
 					m.metas[index] = m.metas[index + 2] - probe_inc
@@ -551,7 +670,7 @@ pub fn (mut m map) delete_1(key voidptr) {
 			m.key_values.delete(kv_index)
 			unsafe {
 				m.metas[index] = 0
-				m.free_key(pkey)
+				m.free_fn(pkey)
 				// Mark key as deleted
 				C.memset(pkey, 0, m.key_bytes)
 			}
@@ -580,7 +699,7 @@ pub fn (m &map) keys() []string {
 		}
 		unsafe {
 			pkey := m.key_values.key(i)
-			m.key_values.clone_key(item, pkey)
+			m.clone_fn(item, pkey)
 			item += m.key_bytes
 		}
 	}
@@ -595,7 +714,7 @@ pub fn (m &map) keys_1() array {
 		for i := 0; i < m.key_values.len; i++ {
 			unsafe {
 				pkey := m.key_values.key(i)
-				m.key_values.clone_key(item, pkey)
+				m.clone_fn(item, pkey)
 				item += m.key_bytes
 			}
 		}
@@ -607,13 +726,14 @@ pub fn (m &map) keys_1() array {
 		}
 		unsafe {
 			pkey := m.key_values.key(i)
-			m.key_values.clone_key(item, pkey)
+			m.clone_fn(item, pkey)
 			item += m.key_bytes
 		}
 	}
 	return keys
 }
 
+// warning: only copies keys, does not clone
 [unsafe]
 pub fn (d &DenseArray) clone() DenseArray {
 	res := DenseArray{
@@ -632,7 +752,6 @@ pub fn (d &DenseArray) clone() DenseArray {
 		}
 		res.data = memdup(d.data, d.cap * d.slot_bytes)
 	}
-	// FIXME clone each key
 	return res
 }
 
@@ -649,8 +768,23 @@ pub fn (m &map) clone() map {
 		metas: &u32(malloc(metasize))
 		extra_metas: m.extra_metas
 		len: m.len
+		has_string_keys: m.has_string_keys
+		hash_fn: m.hash_fn
+		key_eq_fn: m.key_eq_fn
+		clone_fn: m.clone_fn
+		free_fn: m.free_fn
 	}
 	unsafe { C.memcpy(res.metas, m.metas, metasize) }
+	if !m.has_string_keys {
+		return res
+	}
+	// clone keys
+	for i in 0 .. m.key_values.len {
+		if !m.key_values.has_index(i) {
+			continue
+		}
+		m.clone_fn(res.key_values.key(i), m.key_values.key(i))
+	}
 	return res
 }
 
@@ -661,7 +795,7 @@ pub fn (m &map) free() {
 		for i := 0; i < m.key_values.len; i++ {
 			unsafe {
 				pkey := m.key_values.key(i)
-				m.free_key(pkey)
+				m.free_fn(pkey)
 			}
 		}
 	} else {
@@ -671,7 +805,7 @@ pub fn (m &map) free() {
 			}
 			unsafe {
 				pkey := m.key_values.key(i)
-				m.free_key(pkey)
+				m.free_fn(pkey)
 			}
 		}
 		unsafe { free(m.key_values.all_deleted) }

vlib/builtin/map_test.v

@@ -463,3 +463,35 @@ fn test_map_or() {
 	_ = m
 	// num := m['first'] or { return }
 }
+
+fn test_int_keys() {
+	mut m := map[int]int
+	m[3] = 9
+	m[4] = 16
+	assert m.len == 2
+	assert m[3] == 9
+	assert m[4] == 16
+	m[5] += 24
+	m[5]++
+	assert m[5] == 25
+
+	mc := m.clone()
+	assert mc.len == 3
+	mut all := []int{}
+	for k, v in mc {
+		assert m[k] == v
+		all << k
+		all << v
+	}
+	assert all == [3,9,4,16,5,25]
+}
+
+fn test_voidptr_keys() {
+	mut m := map[voidptr]string
+	v := 5
+	m[&v] = 'var'
+	m[&m] = 'map'
+	assert m[&v] == 'var'
+	assert m[&m] == 'map'
+	assert m.len == 2
+}

vlib/hash/wyhash.c.v

@@ -3,9 +3,14 @@ module hash
 //#flag -I @VROOT/thirdparty/wyhash
 //#include "wyhash.h"
 fn C.wyhash(byteptr, u64, u64) u64
+fn C.wyhash64(u64, u64) u64
 
 [inline]
 pub fn wyhash_c(key byteptr, len u64, seed u64) u64 {
 	return C.wyhash(key, len, seed)
 }
 
+[inline]
+pub fn wyhash64_c(a u64, b u64) u64 {
+	return C.wyhash64(a, b)
+}

vlib/v/checker/checker.v

@@ -4433,9 +4433,14 @@ pub fn (mut c Checker) index_expr(mut node ast.IndexExpr) table.Type {
 		return typ.set_nr_muls(0)
 	} else { // [1]
 		index_type := c.expr(node.index)
-		c.check_index_type(typ_sym, index_type, node.pos)
-		if typ_sym.kind == .map && index_type.idx() != table.string_type_idx {
-			c.error('non-string map index (map type `$typ_sym.name`)', node.pos)
+		if typ_sym.kind == .map {
+			info := typ_sym.info as table.Map
+			if !c.check_types(index_type, info.key_type) {
+				err := c.expected_msg(index_type, info.key_type)
+				c.error('invalid key: $err', node.pos)
+			}
+		} else {
+			c.check_index_type(typ_sym, index_type, node.pos)
 		}
 		value_type := c.table.value_type(typ)
 		if value_type != table.void_type {

vlib/v/checker/tests/map_ops.out

@@ -0,0 +1,26 @@
+vlib/v/checker/tests/map_ops.vv:2:18: warning: non-string keys are work in progress
+    1 | fn test_map() {
+    2 |     mut m := map[int]string
+      |                     ^
+    3 |     _ = m[`!`]
+    4 |     m['hi'] = 8
+vlib/v/checker/tests/map_ops.vv:3:7: error: invalid key: expected `int`, not `rune`
+    1 | fn test_map() {
+    2 |     mut m := map[int]string
+    3 |     _ = m[`!`]
+      |          ~~~~~
+    4 |     m['hi'] = 8
+    5 |     m[&m] += 4
+vlib/v/checker/tests/map_ops.vv:4:3: error: invalid key: expected `int`, not `string`
+    2 |     mut m := map[int]string
+    3 |     _ = m[`!`]
+    4 |     m['hi'] = 8
+      |      ~~~~~~
+    5 |     m[&m] += 4
+    6 | }
+vlib/v/checker/tests/map_ops.vv:5:3: error: invalid key: expected `int`, not `&map[int]string`
+    3 |     _ = m[`!`]
+    4 |     m['hi'] = 8
+    5 |     m[&m] += 4
+      |      ~~~~
+    6 | }

vlib/v/checker/tests/map_ops.vv

@@ -0,0 +1,6 @@
+fn test_map() {
+	mut m := map[int]string
+	_ = m[`!`]
+	m['hi'] = 8
+	m[&m] += 4
+}

vlib/v/gen/cgen.v

@@ -1726,6 +1726,7 @@ fn (mut g Gen) gen_assign_stmt(assign_stmt ast.AssignStmt) {
 						g.writeln(');')
 					} else if sym.kind == .map {
 						info := sym.info as table.Map
+						skeytyp := g.typ(info.key_type)
 						styp := g.typ(info.value_type)
 						zero := g.type_default(info.value_type)
 						val_typ := g.table.get_type_symbol(info.value_type)
@@ -1744,7 +1745,7 @@ fn (mut g Gen) gen_assign_stmt(assign_stmt ast.AssignStmt) {
 						} else {
 							g.expr(left.left)
 						}
-						g.write(', &(string[]){')
+						g.write(', &($skeytyp[]){')
 						g.expr(left.index)
 						g.write('}')
 						if val_typ.kind == .function {
@@ -3884,6 +3885,7 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) {
 				}
 			} else if sym.kind == .map {
 				info := sym.info as table.Map
+				key_type_str := g.typ(info.key_type)
 				elem_type_str := g.typ(info.value_type)
 				elem_typ := g.table.get_type_symbol(info.value_type)
 				get_and_set_types := elem_typ.kind in [.struct_, .map]
@@ -3905,7 +3907,7 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) {
 							g.write('.val')
 						}
 					}
-					g.write(', &(string[]){')
+					g.write(', &($key_type_str[]){')
 					g.expr(node.index)
 					g.write('}')
 					if elem_typ.kind == .function {
@@ -3925,7 +3927,7 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) {
 						g.write('&')
 					}
 					g.expr(node.left)
-					g.write(', &(string[]){')
+					g.write(', &($key_type_str[]){')
 					g.expr(node.index)
 					g.write('}, &($elem_type_str[]){ $zero }))')
 				} else {
@@ -3955,7 +3957,7 @@ fn (mut g Gen) index_expr(node ast.IndexExpr) {
 							g.write('.val')
 						}
 					}
-					g.write('), &(string[]){')
+					g.write('), &($key_type_str[]){')
 					g.expr(node.index)
 					g.write('}')
 					if g.is_fn_index_call {

vlib/v/parser/parse_type.v

@@ -52,12 +52,16 @@ pub fn (mut p Parser) parse_map_type() table.Type {
 		// error is reported in parse_type
 		return 0
 	}
-	// key_type_sym := p.get_type_symbol(key_type)
-	// if key_type_sym.kind != .string {
-	if key_type.idx() != table.string_type_idx {
-		p.error('maps can only have string keys for now')
+	if !(key_type in [table.string_type_idx, table.voidptr_type_idx] ||
+		(key_type.is_int() && !key_type.is_ptr())) {
+		s := p.table.type_to_str(key_type)
+		p.error_with_pos('maps only support string, integer, rune or voidptr keys for now (not `$s`)',
+			p.tok.position())
 		return 0
 	}
+	if key_type != table.string_type_idx {
+		p.warn_with_pos('non-string keys are work in progress', p.tok.position())
+	}
 	p.check(.rsbr)
 	value_type := p.parse_type()
 	if value_type.idx() == 0 {
@@ -117,7 +121,10 @@ pub fn (mut p Parser) parse_fn_type(name string) table.Type {
 		is_variadic: is_variadic
 		return_type: return_type
 	}
-	idx := p.table.find_or_register_fn_type(p.mod, func, false, false)
+	// MapFooFn typedefs are manually added in cheaders.v 
+	// because typedefs get generated after the map struct is generated
+	has_decl := p.builtin_mod && name.starts_with('Map') && name.ends_with('Fn')
+	idx := p.table.find_or_register_fn_type(p.mod, func, false, has_decl)
 	return table.new_type(idx)
 }
 

vlib/v/parser/parser.v

@@ -2053,7 +2053,7 @@ fn (mut p Parser) type_decl() ast.TypeDecl {
 	mut type_pos := p.tok.position()
 	mut comments := []ast.Comment{}
 	if p.tok.kind == .key_fn {
-		// function type: `type mycallback fn(string, int)`
+		// function type: `type mycallback = fn(string, int)`
 		fn_name := p.prepend_mod(name)
 		fn_type := p.parse_fn_type(fn_name)
 		comments = p.eat_line_end_comments()