compiler: support comma-ok in map lookup

5 files changed, 33 insertions, 15 deletions

diff --git a/compiler/compiler.go b/compiler/compiler.go
index 4d8330615..964d7763f 100644
--- a/compiler/compiler.go
+++ b/compiler/compiler.go
@@ -2242,9 +2242,6 @@ func (c *Compiler) parseExpr(frame *Frame, expr ssa.Value) (llvm.Value, error) {
 			panic("unreachable")
 		}
 	case *ssa.Lookup:
-		if expr.CommaOk {
-			return llvm.Value{}, errors.New("todo: lookup with comma-ok")
-		}
 		value, err := c.parseExpr(frame, expr.X)
 		if err != nil {
 			return llvm.Value{}, nil
@@ -2273,7 +2270,11 @@ func (c *Compiler) parseExpr(frame *Frame, expr ssa.Value) (llvm.Value, error) {
 			bufPtr := c.builder.CreateGEP(buf, []llvm.Value{index}, "")
 			return c.builder.CreateLoad(bufPtr, ""), nil
 		case *types.Map:
-			return c.emitMapLookup(xType.Key(), expr.Type(), value, index)
+			valueType := expr.Type()
+			if expr.CommaOk {
+				valueType = valueType.(*types.Tuple).At(0).Type()
+			}
+			return c.emitMapLookup(xType.Key(), valueType, value, index, expr.CommaOk)
 		default:
 			panic("unknown lookup type: " + expr.String())
 		}
diff --git a/compiler/map.go b/compiler/map.go
index 23d216da4..d4d318bcc 100644
--- a/compiler/map.go
+++ b/compiler/map.go
@@ -9,29 +9,37 @@ import (
 	"github.com/aykevl/go-llvm"
 )
 
-func (c *Compiler) emitMapLookup(keyType, valueType types.Type, m, key llvm.Value) (llvm.Value, error) {
+func (c *Compiler) emitMapLookup(keyType, valueType types.Type, m, key llvm.Value, commaOk bool) (llvm.Value, error) {
 	llvmValueType, err := c.getLLVMType(valueType)
 	if err != nil {
 		return llvm.Value{}, err
 	}
 	mapValueAlloca := c.builder.CreateAlloca(llvmValueType, "hashmap.value")
 	mapValuePtr := c.builder.CreateBitCast(mapValueAlloca, c.i8ptrType, "hashmap.valueptr")
+	var commaOkValue llvm.Value
 	if t, ok := keyType.(*types.Basic); ok && t.Info()&types.IsString != 0 {
 		// key is a string
 		params := []llvm.Value{m, key, mapValuePtr}
-		c.createRuntimeCall("hashmapStringGet", params, "")
-		return c.builder.CreateLoad(mapValueAlloca, ""), nil
+		commaOkValue = c.createRuntimeCall("hashmapStringGet", params, "")
 	} else if hashmapIsBinaryKey(keyType) {
 		// key can be compared with runtime.memequal
 		keyAlloca := c.builder.CreateAlloca(key.Type(), "hashmap.key")
 		c.builder.CreateStore(key, keyAlloca)
 		keyPtr := c.builder.CreateBitCast(keyAlloca, c.i8ptrType, "hashmap.keyptr")
 		params := []llvm.Value{m, keyPtr, mapValuePtr}
-		c.createRuntimeCall("hashmapBinaryGet", params, "")
-		return c.builder.CreateLoad(mapValueAlloca, ""), nil
+		commaOkValue = c.createRuntimeCall("hashmapBinaryGet", params, "")
 	} else {
 		return llvm.Value{}, errors.New("todo: map lookup key type: " + keyType.String())
 	}
+	mapValue := c.builder.CreateLoad(mapValueAlloca, "")
+	if commaOk {
+		tuple := llvm.Undef(c.ctx.StructType([]llvm.Type{llvmValueType, c.ctx.Int1Type()}, false))
+		tuple = c.builder.CreateInsertValue(tuple, mapValue, 0, "")
+		tuple = c.builder.CreateInsertValue(tuple, commaOkValue, 1, "")
+		return tuple, nil
+	} else {
+		return mapValue, nil
+	}
 }
 
 func (c *Compiler) emitMapUpdate(keyType types.Type, m, key, value llvm.Value) error {
diff --git a/src/runtime/hashmap.go b/src/runtime/hashmap.go
index 9e93928aa..21ef1dd4c 100644
--- a/src/runtime/hashmap.go
+++ b/src/runtime/hashmap.go
@@ -122,7 +122,7 @@ func hashmapSet(m *hashmap, key unsafe.Pointer, value unsafe.Pointer, hash uint3
 
 // Get the value of a specified key, or zero the value if not found.
 //go:nobounds
-func hashmapGet(m *hashmap, key unsafe.Pointer, value unsafe.Pointer, hash uint32, keyEqual func(x, y unsafe.Pointer, n uintptr) bool) {
+func hashmapGet(m *hashmap, key unsafe.Pointer, value unsafe.Pointer, hash uint32, keyEqual func(x, y unsafe.Pointer, n uintptr) bool) bool {
 	numBuckets := uintptr(1) << m.bucketBits
 	bucketNumber := (uintptr(hash) & (numBuckets - 1))
 	bucketSize := unsafe.Sizeof(hashmapBucket{}) + uintptr(m.keySize)*8 + uintptr(m.valueSize)*8
@@ -147,7 +147,7 @@ func hashmapGet(m *hashmap, key unsafe.Pointer, value unsafe.Pointer, hash uint3
 				if keyEqual(key, slotKey, uintptr(m.keySize)) {
 					// Found the key, copy it.
 					memcpy(value, slotValue, uintptr(m.valueSize))
-					return
+					return true
 				}
 			}
 		}
@@ -156,6 +156,7 @@ func hashmapGet(m *hashmap, key unsafe.Pointer, value unsafe.Pointer, hash uint3
 
 	// Did not find the key.
 	memzero(value, uintptr(m.valueSize))
+	return false
 }
 
 // Delete a given key from the map. No-op when the key does not exist in the
@@ -239,9 +240,9 @@ func hashmapBinarySet(m *hashmap, key, value unsafe.Pointer) {
 	hashmapSet(m, key, value, hash, memequal)
 }
 
-func hashmapBinaryGet(m *hashmap, key, value unsafe.Pointer) {
+func hashmapBinaryGet(m *hashmap, key, value unsafe.Pointer) bool {
 	hash := hashmapHash(key, uintptr(m.keySize))
-	hashmapGet(m, key, value, hash, memequal)
+	return hashmapGet(m, key, value, hash, memequal)
 }
 
 func hashmapBinaryDelete(m *hashmap, key unsafe.Pointer) {
@@ -265,9 +266,9 @@ func hashmapStringSet(m *hashmap, key string, value unsafe.Pointer) {
 	hashmapSet(m, unsafe.Pointer(&key), value, hash, hashmapStringEqual)
 }
 
-func hashmapStringGet(m *hashmap, key string, value unsafe.Pointer) {
+func hashmapStringGet(m *hashmap, key string, value unsafe.Pointer) bool {
 	hash := hashmapStringHash(key)
-	hashmapGet(m, unsafe.Pointer(&key), value, hash, hashmapStringEqual)
+	return hashmapGet(m, unsafe.Pointer(&key), value, hash, hashmapStringEqual)
 }
 
 func hashmapStringDelete(m *hashmap, key string) {
diff --git a/testdata/map.go b/testdata/map.go
index fe070625e..9c704cfcf 100644
--- a/testdata/map.go
+++ b/testdata/map.go
@@ -24,6 +24,8 @@ func main() {
 	readMap(testmap2, "ten")
 	delete(testmap2, "six")
 	readMap(testmap2, "seven")
+	lookup(testmap2, "eight")
+	lookup(testmap2, "nokey")
 }
 
 func readMap(m map[string]int, key string) {
@@ -33,3 +35,7 @@ func readMap(m map[string]int, key string) {
 		println(" ", k, "=", v)
 	}
 }
+func lookup(m map[string]int, key string) {
+	value, ok := m[key]
+	println("lookup with comma-ok:", key, value, ok)
+}
diff --git a/testdata/map.txt b/testdata/map.txt
index 4f96c74f1..d076a595a 100644
--- a/testdata/map.txt
+++ b/testdata/map.txt
@@ -46,3 +46,5 @@ map read: seven = 7
   ten = 10
   eleven = 11
   twelve = 12
+lookup with comma-ok: eight 8 true
+lookup with comma-ok: nokey 0 false