From e8ab26875885767ddbcff45290f5ba9d5e909c48 Mon Sep 17 00:00:00 2001
From: Mark Spanbroek <mark@spanbroek.net>
Date: Tue, 4 May 2021 15:55:05 +0200
Subject: [PATCH] Ensure that .? chains work in generic code

---
 questionable/chaining.nim    | 16 ++++++++++++++++
 testmodules/options/test.nim | 17 +++++++++++++++++
 testmodules/result/test.nim  | 17 +++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/questionable/chaining.nim b/questionable/chaining.nim
index bee253d..7c8ca14 100644
--- a/questionable/chaining.nim
+++ b/questionable/chaining.nim
@@ -1,6 +1,12 @@
 import std/options
 import std/macros
 
+func isSym(node: NimNode): bool =
+  node.kind in {nnkSym, nnkOpenSymChoice, nnkClosedSymChoice}
+
+func expectSym(node: NimNode) =
+  node.expectKind({nnkSym, nnkOpenSymChoice, nnkClosedSymChoice})
+
 template `.?`*(option: typed, identifier: untyped{nkIdent}): untyped =
   # chain is of shape: option.?identifier
   option ->? option.unsafeGet.identifier
@@ -34,7 +40,17 @@ macro `.?`*(option: typed, call: untyped{nkCall}): untyped =
     call[0] = right
     call.insert(1, quote do: `option`.?`left`)
     call
+  elif procedure.isSym and $procedure == "[]":
+    # chain is of shape: option.?left[right] after semantic analysis
+    let left = call[1]
+    call[1] = quote do: `option`.?`left`
+    call
   else:
     # chain is of shape: option.?procedure(arguments)
     call.insert(1, quote do: `option`.unsafeGet)
     quote do: `option` ->? `call`
+
+macro `.?`*(option: typed, symbol: untyped): untyped =
+  symbol.expectSym()
+  let expression = ident($symbol)
+  quote do: `option`.?`expression`
diff --git a/testmodules/options/test.nim b/testmodules/options/test.nim
index d6b0087..260fff3 100644
--- a/testmodules/options/test.nim
+++ b/testmodules/options/test.nim
@@ -38,6 +38,23 @@ suite "optionals":
     check a.?deduplicate()[0] + 1 == 42.some
     check a.?deduplicate.map(x => x) == @[41, 42].some
 
+  test ".? chains work in generic code":
+    proc test[T](a: ?T) =
+      check a.?len == 2.some
+      check a.?len.?uint8 == 2'u8.some
+      check a.?len() == 2.some
+      check a.?distribute(2).?len() == 2.some
+      check a.?len.unsafeGet == 2
+      check a.?len.unsafeGet.uint8.uint64 == 2'u64
+      check a.?len.unsafeGet() == 2
+      check a.?len.unsafeGet().uint8.uint64 == 2'u64
+      check a.?deduplicate()[0].?uint8.?uint64 == 41'u64.some
+      check a.?len + 1 == 3.some
+      check a.?deduplicate()[0] + 1 == 42.some
+      check a.?deduplicate.map(x => x) == @[41, 42].some
+
+    test @[41, 42].some
+
   test "[] can be used for indexing optionals":
     let a: ?seq[int] = @[1, 2, 3].some
     let b: ?seq[int] = seq[int].none
diff --git a/testmodules/result/test.nim b/testmodules/result/test.nim
index 9ae49cc..0708ca5 100644
--- a/testmodules/result/test.nim
+++ b/testmodules/result/test.nim
@@ -41,6 +41,23 @@ suite "result":
     check (a.?deduplicate()[0] + 1 == 42.success)
     check (a.?deduplicate.map(x => x) == @[41, 42].success)
 
+  test ".? chains work in generic code":
+    proc test[T](a: ?!T) =
+      check (a.?len == 2.success)
+      check (a.?len.?uint8 == 2'u8.success)
+      check (a.?len() == 2.success)
+      check (a.?distribute(2).?len() == 2.success)
+      check (a.?len.unsafeGet == 2)
+      check (a.?len.unsafeGet.uint8.uint64 == 2'u64)
+      check (a.?len.unsafeGet() == 2)
+      check (a.?len.unsafeGet().uint8.uint64 == 2'u64)
+      check (a.?deduplicate()[0].?uint8.?uint64 == 41'u64.success)
+      check (a.?len + 1 == 3.success)
+      check (a.?deduplicate()[0] + 1 == 42.success)
+      check (a.?deduplicate.map(x => x) == @[41, 42].success)
+
+    test @[41, 42].success
+
   test "[] can be used for indexing results":
     let a: ?!seq[int] = @[1, 2, 3].success
     let b: ?!seq[int] = seq[int].failure error