From e374311f70f43e9786de3f2ca3e6cd3bca5aaedd Mon Sep 17 00:00:00 2001
From: jangko <jangko128@gmail.com>
Date: Fri, 7 Jul 2023 13:22:22 +0700
Subject: [PATCH] add custom literal feature

---
 stint.nim               | 12 ++++++++++++
 stint/io.nim            | 26 +++++++++++++++++++++++++-
 tests/all_tests.nim     |  4 +++-
 tests/test_features.nim | 29 +++++++++++++++++++++++++++++
 4 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_features.nim

diff --git a/stint.nim b/stint.nim
index a3ba9ae..1fde323 100644
--- a/stint.nim
+++ b/stint.nim
@@ -31,3 +31,15 @@ func i256*(s: string): Int256 {.inline.} = s.parse(Int256)
 # According to nim manual, you can write something like 1234567890'u256
 # or 1234567890'i256, and the number will be passed as string to the
 # constructor
+
+func `'i128`*(s: static string): Int128 {.inline.} = 
+  customLiteral(Int128, s)
+
+func `'i256`*(s: static string): Int256 {.inline.} = 
+  customLiteral(Int256, s)
+  
+func `'u128`*(s: static string): UInt128 {.inline.} = 
+  customLiteral(UInt128, s)
+
+func `'u256`*(s: static string): UInt256 {.inline.} = 
+  customLiteral(UInt256, s)
diff --git a/stint/io.nim b/stint/io.nim
index ba2362f..a0ea092 100644
--- a/stint/io.nim
+++ b/stint/io.nim
@@ -261,7 +261,7 @@ func parse*[bits: static[int]](input: string,
 
 func parse*[bits: static[int]](input: string,
                                T: typedesc[StInt[bits]],
-                               radix: static[int8] = 10): T =
+                               radix: static[uint8] = 10): T =
   ## Parse a string and store the result in a Stint[bits] or StUint[bits].
 
   static: doAssert (radix >= 2) and radix <= 16, "Only base from 2..16 are supported"
@@ -465,3 +465,27 @@ template toBytesBE*[bits: static[int]](n: StInt[bits]): array[bits div 8, byte]
   result = n.impl.toBytesBE()
 
 {.pop.}
+
+func getRadix(s: static string): uint8 {.compileTime.} =
+  if s.len <= 2:
+    return 10
+  
+  # maybe have prefix have prefix
+  if s[0] != '0':
+    return 10
+
+  if s[1] == 'b':
+    return 2
+      
+  if s[1] == 'o':
+    return 8
+    
+  if s[1] == 'x':
+    return 16
+  
+func customLiteral*(T: type SomeBigInteger, s: static string): T =
+  when s.len == 0:
+    doAssert(false, "customLiteral cannot accept param with zero length")
+    
+  const radix = getRadix(s)
+  parse(s, T, radix)
diff --git a/tests/all_tests.nim b/tests/all_tests.nim
index 8fb0fd4..cc6a34d 100644
--- a/tests/all_tests.nim
+++ b/tests/all_tests.nim
@@ -35,4 +35,6 @@ import
   test_io,
   test_conversion,
   t_randomized_divmod,
-  test_bugfix
+  test_bugfix,
+  test_features
+  
\ No newline at end of file
diff --git a/tests/test_features.nim b/tests/test_features.nim
new file mode 100644
index 0000000..814e87f
--- /dev/null
+++ b/tests/test_features.nim
@@ -0,0 +1,29 @@
+# Copyright 2023 Status Research & Development GmbH
+# Licensed under either of
+#
+#  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or http://www.apache.org/licenses/LICENSE-2.0)
+#  * MIT license ([LICENSE-MIT](LICENSE-MIT) or http://opensource.org/licenses/MIT)
+#
+# at your option. This file may not be copied, modified, or distributed except according to those terms.
+
+import ../stint, unittest
+
+suite "new features":
+  test "custom literal":
+    const
+      a = 0xabcdef0123456'u128
+      b = 0xabcdef0123456'u256
+      c = -100'i128
+      d = -50000'i256
+
+    let
+      x = 0b111100011'u128
+      y = 0o777766666'u256
+
+    check:
+      a == 0xabcdef0123456.u128
+      b == 0xabcdef0123456.u256
+      c == -100.i128
+      d == -50000.i256
+      x == 0b111100011.u128
+      y == 0o777766666.u256