add custom pragma support for var and let symbols (#9582)

* add custom pragma support for var and let symbols
* updated changelog for custom pragmas on var and let symbols
* add oldast switch for backwards compatibility

9 files changed, 77 insertions, 20 deletions

diff --git a/changelog.md b/changelog.md
index b3ccc1b2e..05f65516d 100644
--- a/changelog.md
+++ b/changelog.md
@@ -23,7 +23,8 @@
 - The undocumented ``#? strongSpaces`` parsing mode has been removed.
 - The `not` operator is now always a unary operator, this means that code like
   ``assert not isFalse(3)`` compiles.
-
+- `getImpl` on a `var` or `let` symbol will now return the full `IdentDefs`
+  tree from the symbol declaration instead of just the initializer portion.
 
 #### Breaking changes in the standard library
 
@@ -133,8 +134,9 @@ proc enumToString*(enums: openArray[enum]): string =
   the `gcsafe` pragma block.
 - added os.getCurrentProcessId()
 - User defined pragmas are now allowed in the pragma blocks
-- Pragma blocks are now longer eliminated from the typed AST tree to preserve
+- Pragma blocks are no longer eliminated from the typed AST tree to preserve
   pragmas for further analysis by macros
+- Custom pragmas are now supported for `var` and `let` symbols.
 
 ### Language changes
 
@@ -143,10 +145,10 @@ proc enumToString*(enums: openArray[enum]): string =
   it's more recognizable and allows tools like github to recognize it as Nim,
   see [#9647](https://github.com/nim-lang/Nim/issues/9647).
   The previous extension will continue to work.
-- Pragma syntax is now consistent. Previous syntax where type pragmas did not 
+- Pragma syntax is now consistent. Previous syntax where type pragmas did not
   follow the type name is now deprecated. Also pragma before generic parameter
   list is deprecated to be consistent with how pragmas are used with a proc. See
-  [#8514](https://github.com/nim-lang/Nim/issues/8514) and 
+  [#8514](https://github.com/nim-lang/Nim/issues/8514) and
   [#1872](https://github.com/nim-lang/Nim/issues/1872) for further details.
 
 
diff --git a/compiler/ast.nim b/compiler/ast.nim
index 5f5f296cb..24891d6d3 100644
--- a/compiler/ast.nim
+++ b/compiler/ast.nim
@@ -1087,6 +1087,13 @@ proc newSym*(symKind: TSymKind, name: PIdent, owner: PSym,
   when debugIds:
     registerId(result)
 
+proc astdef*(s: PSym): PNode =
+  # get only the definition (initializer) portion of the ast
+  if s.ast != nil and s.ast.kind == nkIdentDefs:
+    s.ast[2]
+  else:
+    s.ast
+
 proc isMetaType*(t: PType): bool =
   return t.kind in tyMetaTypes or
          (t.kind == tyStatic and t.n == nil) or
diff --git a/compiler/commands.nim b/compiler/commands.nim
index 8b73884e8..47687046f 100644
--- a/compiler/commands.nim
+++ b/compiler/commands.nim
@@ -291,6 +291,7 @@ proc testCompileOption*(conf: ConfigRef; switch: string, info: TLineInfo): bool
   of "patterns": result = contains(conf.options, optPatterns)
   of "excessivestacktrace": result = contains(conf.globalOptions, optExcessiveStackTrace)
   of "nilseqs": result = contains(conf.options, optNilSeqs)
+  of "oldast": result = contains(conf.options, optOldAst)
   else: invalidCmdLineOption(conf, passCmd1, switch, info)
 
 proc processPath(conf: ConfigRef; path: string, info: TLineInfo,
@@ -508,6 +509,7 @@ proc processSwitch*(switch, arg: string, pass: TCmdLinePass, info: TLineInfo;
       localError(conf, info, errOnOrOffExpectedButXFound % arg)
   of "laxstrings": processOnOffSwitch(conf, {optLaxStrings}, arg, pass, info)
   of "nilseqs": processOnOffSwitch(conf, {optNilSeqs}, arg, pass, info)
+  of "oldast": processOnOffSwitch(conf, {optOldAst}, arg, pass, info)
   of "checks", "x": processOnOffSwitch(conf, ChecksOptions, arg, pass, info)
   of "floatchecks":
     processOnOffSwitch(conf, {optNaNCheck, optInfCheck}, arg, pass, info)
diff --git a/compiler/guards.nim b/compiler/guards.nim
index a01c023e4..46e18d3bf 100644
--- a/compiler/guards.nim
+++ b/compiler/guards.nim
@@ -257,9 +257,9 @@ proc canon*(n: PNode; o: Operators): PNode =
     for i in 0 ..< n.len:
       result.sons[i] = canon(n.sons[i], o)
   elif n.kind == nkSym and n.sym.kind == skLet and
-      n.sym.ast.getMagic in (someEq + someAdd + someMul + someMin +
+      n.sym.astdef.getMagic in (someEq + someAdd + someMul + someMin +
       someMax + someHigh + {mUnaryLt} + someSub + someLen + someDiv):
-    result = n.sym.ast.copyTree
+    result = n.sym.astdef.copyTree
   else:
     result = n
   case result.getMagic
@@ -395,8 +395,8 @@ proc usefulFact(n: PNode; o: Operators): PNode =
     #   if a:
     #     ...
     # We make can easily replace 'a' by '2 < x' here:
-    if n.sym.ast != nil:
-      result = usefulFact(n.sym.ast, o)
+    if n.sym.astdef != nil:
+      result = usefulFact(n.sym.astdef, o)
   elif n.kind == nkStmtListExpr:
     result = usefulFact(n.lastSon, o)
 
diff --git a/compiler/options.nim b/compiler/options.nim
index 3730d16f8..ea159ee1d 100644
--- a/compiler/options.nim
+++ b/compiler/options.nim
@@ -40,7 +40,8 @@ type                          # please make sure we have under 32 options
     optMemTracker,
     optHotCodeReloading,
     optLaxStrings,
-    optNilSeqs
+    optNilSeqs,
+    optOldAst
 
   TOptions* = set[TOption]
   TGlobalOption* = enum       # **keep binary compatible**
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index 39f200ba8..12283e042 100644
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -330,9 +330,9 @@ proc semIdentDef(c: PContext, n: PNode, kind: TSymKind): PSym =
 proc checkNilable(c: PContext; v: PSym) =
   if {sfGlobal, sfImportC} * v.flags == {sfGlobal} and
       {tfNotNil, tfNeedsInit} * v.typ.flags != {}:
-    if v.ast.isNil:
+    if v.astdef.isNil:
       message(c.config, v.info, warnProveInit, v.name.s)
-    elif tfNotNil in v.typ.flags and tfNotNil notin v.ast.typ.flags:
+    elif tfNotNil in v.typ.flags and tfNotNil notin v.astdef.typ.flags:
       message(c.config, v.info, warnProveInit, v.name.s)
 
 include semasgn
@@ -518,8 +518,6 @@ proc semVarOrLet(c: PContext, n: PNode, symkind: TSymKind): PNode =
               message(c.config, a.info, warnShadowIdent, v.name.s)
       if a.kind != nkVarTuple:
         if def.kind != nkEmpty:
-          # this is needed for the evaluation pass and for the guard checking:
-          v.ast = def
           if sfThread in v.flags: localError(c.config, def.info, errThreadvarCannotInit)
         setVarType(c, v, typ)
         b = newNodeI(nkIdentDefs, a.info)
@@ -531,6 +529,23 @@ proc semVarOrLet(c: PContext, n: PNode, symkind: TSymKind): PNode =
         addSon(b, a.sons[length-2])
         addSon(b, copyTree(def))
         addToVarSection(c, result, n, b)
+        if optOldAst in c.config.options:
+          if def.kind != nkEmpty:
+            v.ast = def
+        else:
+          # this is needed for the evaluation pass, guard checking
+          #  and custom pragmas:
+          var ast = newNodeI(nkIdentDefs, a.info)
+          if a[j].kind == nkPragmaExpr:
+            var p = newNodeI(nkPragmaExpr, a.info)
+            p.add newSymNode(v)
+            p.add a[j][1].copyTree
+            ast.add p
+          else:
+            ast.add newSymNode(v)
+          ast.add a.sons[length-2].copyTree
+          ast.add def
+          v.ast = ast
       else:
         if def.kind in {nkPar, nkTupleConstr}: v.ast = def[j]
         # bug #7663, for 'nim check' this can be a non-tuple:
diff --git a/doc/advopt.txt b/doc/advopt.txt
index 5ddd064ef..2e91deed9 100644
--- a/doc/advopt.txt
+++ b/doc/advopt.txt
@@ -76,6 +76,7 @@ Advanced options:
                             strings is allowed; only for backwards compatibility
   --nilseqs:on|off          allow 'nil' for strings/seqs for
                             backwards compatibility
+  --oldast:on|off           use old AST for backwards compatibility
   --skipCfg                 do not read the general configuration file
   --skipUserCfg             do not read the user's configuration file
   --skipParentCfg           do not read the parent dirs' configuration files
diff --git a/lib/core/macros.nim b/lib/core/macros.nim
index 291858bed..64334161e 100644
--- a/lib/core/macros.nim
+++ b/lib/core/macros.nim
@@ -1402,8 +1402,14 @@ proc customPragmaNode(n: NimNode): NimNode =
     let impl = n.getImpl()
     if impl.kind in RoutineNodes:
       return impl.pragma
+    elif impl.kind == nnkIdentDefs and impl[0].kind == nnkPragmaExpr:
+      return impl[0][1]
     else:
-      return typ.getImpl()[0][1]
+      let timpl = typ.getImpl()
+      if timpl.len>0 and timpl[0].len>1:
+        return timpl[0][1]
+      else:
+        return timpl
 
   if n.kind in {nnkDotExpr, nnkCheckedFieldExpr}:
     let name = $(if n.kind == nnkCheckedFieldExpr: n[0][1] else: n[1])
diff --git a/tests/pragmas/tcustom_pragma.nim b/tests/pragmas/tcustom_pragma.nim
index 0bc4d2f18..fefcc0b5f 100644
--- a/tests/pragmas/tcustom_pragma.nim
+++ b/tests/pragmas/tcustom_pragma.nim
@@ -175,24 +175,47 @@ var foo: Something
 foo.cardinal = north
 doAssert foo.b.hasCustomPragma(thingy) == true
 
-
-proc myproc(s: string): int = 
+proc myproc(s: string): int =
   {.thingy.}:
     s.len
 
 doAssert myproc("123") == 3
 
 let xx = compiles:
-  proc myproc_bad(s: string): int = 
+  proc myproc_bad(s: string): int =
     {.not_exist.}:
       s.len
 doAssert: xx == false
 
-
-macro checkSym(s: typed{nkSym}): untyped = 
+macro checkSym(s: typed{nkSym}): untyped =
   let body = s.getImpl.body
   doAssert body[1].kind == nnkPragmaBlock
   doAssert body[1][0].kind == nnkPragma
   doAssert body[1][0][0] == bindSym"thingy"
 
-checkSym(myproc)
\ No newline at end of file
+checkSym(myproc)
+
+# var and let pragmas
+block:
+  template myAttr() {.pragma.}
+  template myAttr2(x: int) {.pragma.}
+  template myAttr3(x: string) {.pragma.}
+
+  let a {.myAttr,myAttr2(2),myAttr3:"test".}: int = 0
+  let b {.myAttr,myAttr2(2),myAttr3:"test".} = 0
+  var x {.myAttr,myAttr2(2),myAttr3:"test".}: int = 0
+  var y {.myAttr,myAttr2(2),myAttr3:"test".}: int
+  var z {.myAttr,myAttr2(2),myAttr3:"test".} = 0
+
+  template check(s: untyped) =
+    doAssert s.hasCustomPragma(myAttr)
+    doAssert s.hasCustomPragma(myAttr2)
+    doAssert s.getCustomPragmaVal(myAttr2) == 2
+    doAssert s.hasCustomPragma(myAttr3)
+    doAssert s.getCustomPragmaVal(myAttr3) == "test"
+
+  check(a)
+  check(b)
+  check(x)
+  check(y)
+  check(z)