ORC: support for custom =trace procs (#18459)

* ORC: support custom =trace procs (WIP)
* Update tests/arc/tcustomtrace.nim

Co-authored-by: Clyybber <darkmine956@gmail.com>

* =trace is now documented and seems to work
* make test green

Co-authored-by: Clyybber <darkmine956@gmail.com>

13 files changed, 321 insertions, 11 deletions

diff --git a/changelog.md b/changelog.md
index 849640e06..bc64e9579 100644
--- a/changelog.md
+++ b/changelog.md
@@ -388,6 +388,8 @@
 
 - `typeof(voidStmt)` now works and returns `void`.
 
+- The `gc:orc` algorithm was refined so that custom container types can participate in the
+  cycle collection process.
 
 
 ## Compiler changes
diff --git a/compiler/ast.nim b/compiler/ast.nim
index 26050426a..f96e464af 100644
--- a/compiler/ast.nim
+++ b/compiler/ast.nim
@@ -667,7 +667,7 @@ type
     mIsPartOf, mAstToStr, mParallel,
     mSwap, mIsNil, mArrToSeq,
     mNewString, mNewStringOfCap, mParseBiggestFloat,
-    mMove, mWasMoved, mDestroy,
+    mMove, mWasMoved, mDestroy, mTrace,
     mDefault, mUnown, mIsolate, mAccessEnv, mReset,
     mArray, mOpenArray, mRange, mSet, mSeq, mVarargs,
     mRef, mPtr, mVar, mDistinct, mVoid, mTuple,
diff --git a/compiler/ccgexprs.nim b/compiler/ccgexprs.nim
index c07ddf1db..09326ceeb 100644
--- a/compiler/ccgexprs.nim
+++ b/compiler/ccgexprs.nim
@@ -2422,6 +2422,7 @@ proc genMagicExpr(p: BProc, e: PNode, d: var TLoc, op: TMagic) =
   of mDestroy: genDestroy(p, e)
   of mAccessEnv: unaryExpr(p, e, d, "$1.ClE_0")
   of mSlice: genSlice(p, e, d)
+  of mTrace: discard "no code to generate"
   else:
     when defined(debugMagics):
       echo p.prc.name.s, " ", p.prc.id, " ", p.prc.flags, " ", p.prc.ast[genericParamsPos].kind
diff --git a/compiler/condsyms.nim b/compiler/condsyms.nim
index b3f8465f0..33cd1818b 100644
--- a/compiler/condsyms.nim
+++ b/compiler/condsyms.nim
@@ -136,3 +136,4 @@ proc initDefines*(symbols: StringTableRef) =
   defineSymbol("nimHasTypeofVoid")
   defineSymbol("nimHasDragonBox")
   defineSymbol("nimHasHintAll")
+  defineSymbol("nimHasTrace")
diff --git a/compiler/jsgen.nim b/compiler/jsgen.nim
index a9339c7fd..82fba02a6 100644
--- a/compiler/jsgen.nim
+++ b/compiler/jsgen.nim
@@ -2079,7 +2079,7 @@ proc genMagic(p: PProc, n: PNode, r: var TCompRes) =
       gen(p, n[1], x)
       useMagic(p, "nimCopy")
       r.res = "nimCopy(null, $1, $2)" % [x.rdLoc, genTypeInfo(p, n.typ)]
-  of mDestroy: discard "ignore calls to the default destructor"
+  of mDestroy, mTrace: discard "ignore calls to the default destructor"
   of mOrd: genOrd(p, n, r)
   of mLengthStr, mLengthSeq, mLengthOpenArray, mLengthArray:
     var x: TCompRes
diff --git a/compiler/liftdestructors.nim b/compiler/liftdestructors.nim
index 1824557e6..b2b70edee 100644
--- a/compiler/liftdestructors.nim
+++ b/compiler/liftdestructors.nim
@@ -416,11 +416,23 @@ proc considerUserDefinedOp(c: var TLiftCtx; t: PType; body, x, y: PNode): bool =
       body.add destructorCall(c, op, x)
       result = true
     #result = addDestructorCall(c, t, body, x)
-  of attachedAsgn, attachedSink, attachedTrace:
+  of attachedAsgn, attachedSink:
     var op = getAttachedOp(c.g, t, c.kind)
     result = considerAsgnOrSink(c, t, body, x, y, op)
     if op != nil:
       setAttachedOp(c.g, c.idgen.module, t, c.kind, op)
+  of attachedTrace:
+    var op = getAttachedOp(c.g, t, c.kind)
+    if op != nil and sfOverriden in op.flags:
+      if op.ast.isGenericRoutine:
+        # patch generic =trace:
+        op = instantiateGeneric(c, op, t, t.typeInst)
+        setAttachedOp(c.g, c.idgen.module, t, c.kind, op)
+
+    result = considerAsgnOrSink(c, t, body, x, y, op)
+    if op != nil:
+      setAttachedOp(c.g, c.idgen.module, t, c.kind, op)
+
   of attachedDeepCopy:
     let op = getAttachedOp(c.g, t, attachedDeepCopy)
     if op != nil:
@@ -1065,7 +1077,7 @@ proc createTypeBoundOps(g: ModuleGraph; c: PContext; orig: PType; info: TLineInf
   # 4. We have a custom destructor.
   # 5. We have a (custom) generic destructor.
 
-  # we do not generate '=trace' nor '=dispose' procs if we
+  # we do not generate '=trace' procs if we
   # have the cycle detection disabled, saves code size.
   let lastAttached = if g.config.selectedGC == gcOrc: attachedTrace
                      else: attachedSink
diff --git a/compiler/semmagic.nim b/compiler/semmagic.nim
index d3f26e630..04f4c4729 100644
--- a/compiler/semmagic.nim
+++ b/compiler/semmagic.nim
@@ -551,6 +551,12 @@ proc magicsAfterOverloadResolution(c: PContext, n: PNode,
     let op = getAttachedOp(c.graph, t, attachedDestructor)
     if op != nil:
       result[0] = newSymNode(op)
+  of mTrace:
+    result = n
+    let t = n[1].typ.skipTypes(abstractVar)
+    let op = getAttachedOp(c.graph, t, attachedTrace)
+    if op != nil:
+      result[0] = newSymNode(op)
   of mUnown:
     result = semUnown(c, n)
   of mExists, mForall:
diff --git a/compiler/semstmts.nim b/compiler/semstmts.nim
index c3402aee3..fd2f8a1d9 100644
--- a/compiler/semstmts.nim
+++ b/compiler/semstmts.nim
@@ -1649,6 +1649,8 @@ proc bindTypeHook(c: PContext; s: PSym; n: PNode; op: TTypeAttachedOp) =
   var noError = false
   let cond = if op == attachedDestructor:
                t.len == 2 and t[0] == nil and t[1].kind == tyVar
+             elif op == attachedTrace:
+               t.len == 3 and t[0] == nil and t[1].kind == tyVar and t[2].kind == tyPointer
              else:
                t.len >= 2 and t[0] == nil
 
@@ -1673,8 +1675,12 @@ proc bindTypeHook(c: PContext; s: PSym; n: PNode; op: TTypeAttachedOp) =
         localError(c.config, n.info, errGenerated,
           "type bound operation `" & s.name.s & "` can be defined only in the same module with its type (" & obj.typeToString() & ")")
   if not noError and sfSystemModule notin s.owner.flags:
-    localError(c.config, n.info, errGenerated,
-      "signature for '" & s.name.s & "' must be proc[T: object](x: var T)")
+    if op == attachedTrace:
+      localError(c.config, n.info, errGenerated,
+        "signature for '=trace' must be proc[T: object](x: var T; env: pointer)")
+    else:
+      localError(c.config, n.info, errGenerated,
+        "signature for '" & s.name.s & "' must be proc[T: object](x: var T)")
   incl(s.flags, sfUsed)
   incl(s.flags, sfOverriden)
 
@@ -1752,7 +1758,8 @@ proc semOverride(c: PContext, s: PSym, n: PNode) =
       localError(c.config, n.info, errGenerated,
                 "signature for '" & s.name.s & "' must be proc[T: object](x: var T; y: T)")
   of "=trace":
-    bindTypeHook(c, s, n, attachedTrace)
+    if s.magic != mTrace:
+      bindTypeHook(c, s, n, attachedTrace)
   else:
     if sfOverriden in s.flags:
       localError(c.config, n.info, errGenerated,
diff --git a/compiler/vmgen.nim b/compiler/vmgen.nim
index 14f9e0b30..e5e4a854e 100644
--- a/compiler/vmgen.nim
+++ b/compiler/vmgen.nim
@@ -1362,7 +1362,7 @@ proc genMagic(c: PCtx; n: PNode; dest: var TDest; m: TMagic) =
     globalError(c.config, n.info, sizeOfLikeMsg("offsetof"))
   of mRunnableExamples:
     discard "just ignore any call to runnableExamples"
-  of mDestroy: discard "ignore calls to the default destructor"
+  of mDestroy, mTrace: discard "ignore calls to the default destructor"
   of mMove:
     let arg = n[1]
     let a = c.genx(arg)
diff --git a/doc/destructors.rst b/doc/destructors.rst
index c99606b95..a94ccb9e7 100644
--- a/doc/destructors.rst
+++ b/doc/destructors.rst
@@ -42,6 +42,12 @@ written as:
       for i in 0..<x.len: `=destroy`(x.data[i])
       dealloc(x.data)
 
+  proc `=trace`[T](x: var myseq[T]; env: pointer) =
+    # `=trace` allows the cycle collector `--gc:orc`
+    # to understand how to trace the object graph.
+    if x.data != nil:
+      for i in 0..<x.len: `=trace`(x.data[i], env)
+
   proc `=copy`*[T](a: var myseq[T]; b: myseq[T]) =
     # do nothing for self-assignments:
     if a.data == b.data: return
@@ -198,6 +204,37 @@ that otherwise would lead to a copy is prevented at compile-time. This looks lik
 but a custom error message (e.g., `{.error: "custom error".}`) will not be emitted
 by the compiler. Notice that there is no `=` before the `{.error.}` pragma.
 
+
+`=trace` hook
+---------------
+
+A custom **container** type can support Nim's cycle collector `--gc:orc` via
+the `=trace` hook. If the container does not implement `=trace`, cyclic data
+structure which are constructed with the help of the container might leak
+memory or resources, but memory safety is not compromised.
+
+The prototype of this hook for a type `T` needs to be:
+
+.. code-block:: nim
+
+  proc `=trace`(dest: var T; env: pointer)
+
+`env` is used by ORC to keep track of its internal state, it should be passed around
+to calls of the built-in `=trace` operation.
+
+The general pattern in `=trace` looks like:
+
+.. code-block:: nim
+
+  proc `=trace`(dest: var T; env: pointer) =
+    for child in childrenThatCanContainPointers(dest):
+      `=trace`(child, env)
+
+
+**Note**: The `=trace` hooks is currently more experimental and less refined
+than the other hooks.
+
+
 Move semantics
 ==============
 
diff --git a/doc/manual.rst b/doc/manual.rst
index 9ddf02009..34781141e 100644
--- a/doc/manual.rst
+++ b/doc/manual.rst
@@ -3840,9 +3840,8 @@ the operator is in scope (including if it is private).
   # will still be called upon exiting scope
   doAssert witness == 3
 
-Type bound operators currently include:
-`=destroy`, `=copy`, `=sink`, `=trace`, `=deepcopy`
-(some of which are still implementation defined and not yet documented).
+Type bound operators are:
+`=destroy`, `=copy`, `=sink`, `=trace`, `=deepcopy`.
 
 For more details on some of those procs, see
 `Lifetime-tracking hooks <destructors.html#lifetimeminustracking-hooks>`_.
diff --git a/lib/system.nim b/lib/system.nim
index 090ab309e..8aa9bd718 100644
--- a/lib/system.nim
+++ b/lib/system.nim
@@ -488,6 +488,11 @@ proc `=sink`*[T](x: var T; y: T) {.inline, magic: "Asgn".} =
   ## Generic `sink`:idx: implementation that can be overridden.
   shallowCopy(x, y)
 
+when defined(nimHasTrace):
+  proc `=trace`*[T](x: var T; env: pointer) {.inline, magic: "Trace".} =
+    ## Generic `trace`:idx: implementation that can be overridden.
+    discard
+
 type
   HSlice*[T, U] = object   ## "Heterogeneous" slice type.
     a*: T                  ## The lower bound (inclusive).
diff --git a/tests/arc/tcustomtrace.nim b/tests/arc/tcustomtrace.nim
new file mode 100644
index 000000000..3977194d9
--- /dev/null
+++ b/tests/arc/tcustomtrace.nim
@@ -0,0 +1,240 @@
+discard """
+  outputsub: '''1
+2
+3
+4
+5
+6
+89
+90
+90
+0 0 1
+0 1 2
+0 2 3
+1 0 4
+1 1 5
+1 2 6
+1 3 7
+after 6 6
+MEM 0'''
+joinable: false
+  cmd: "nim c --gc:orc -d:useMalloc $file"
+  valgrind: "true"
+"""
+
+import typetraits
+
+type
+  myseq*[T] = object
+    len, cap: int
+    data: ptr UncheckedArray[T]
+
+# XXX make code memory safe for overflows in '*'
+var
+  allocCount, deallocCount: int
+
+proc `=destroy`*[T](x: var myseq[T]) =
+  if x.data != nil:
+    when not supportsCopyMem(T):
+      for i in 0..<x.len: `=destroy`(x[i])
+    dealloc(x.data)
+    inc deallocCount
+    x.data = nil
+    x.len = 0
+    x.cap = 0
+
+proc `=copy`*[T](a: var myseq[T]; b: myseq[T]) =
+  if a.data == b.data: return
+  if a.data != nil:
+    `=destroy`(a)
+    #dealloc(a.data)
+    #inc deallocCount
+    #a.data = nil
+  a.len = b.len
+  a.cap = b.cap
+  if b.data != nil:
+    a.data = cast[type(a.data)](alloc(a.cap * sizeof(T)))
+    inc allocCount
+    when supportsCopyMem(T):
+      copyMem(a.data, b.data, a.cap * sizeof(T))
+    else:
+      for i in 0..<a.len:
+        a.data[i] = b.data[i]
+
+proc `=sink`*[T](a: var myseq[T]; b: myseq[T]) =
+  if a.data != nil and a.data != b.data:
+    dealloc(a.data)
+    inc deallocCount
+  a.len = b.len
+  a.cap = b.cap
+  a.data = b.data
+
+proc `=trace`*[T](x: var myseq[T]; env: pointer) =
+  if x.data != nil:
+    for i in 0..<x.len: `=trace`(x[i], env)
+
+proc resize[T](s: var myseq[T]) =
+  let oldCap = s.cap
+  if oldCap == 0: s.cap = 8
+  else: s.cap = (s.cap * 3) shr 1
+  if s.data == nil: inc allocCount
+  s.data = cast[typeof(s.data)](realloc0(s.data, oldCap * sizeof(T), s.cap * sizeof(T)))
+
+proc reserveSlot[T](x: var myseq[T]): ptr T =
+  if x.len >= x.cap: resize(x)
+  result = addr(x.data[x.len])
+  inc x.len
+
+template add*[T](x: var myseq[T]; y: T) =
+  reserveSlot(x)[] = y
+
+proc shrink*[T](x: var myseq[T]; newLen: int) =
+  assert newLen <= x.len
+  assert newLen >= 0
+  when not supportsCopyMem(T):
+    for i in countdown(x.len - 1, newLen - 1):
+      `=destroy`(x.data[i])
+  x.len = newLen
+
+proc grow*[T](x: var myseq[T]; newLen: int; value: T) =
+  if newLen <= x.len: return
+  assert newLen >= 0
+  let oldCap = x.cap
+  if oldCap == 0: x.cap = newLen
+  else: x.cap = max(newLen, (oldCap * 3) shr 1)
+  if x.data == nil: inc allocCount
+  x.data = cast[type(x.data)](realloc0(x.data, oldCap * sizeof(T), x.cap * sizeof(T)))
+  for i in x.len..<newLen:
+    x.data[i] = value
+  x.len = newLen
+
+template default[T](t: typedesc[T]): T =
+  var v: T
+  v
+
+proc setLen*[T](x: var myseq[T]; newLen: int) {.deprecated.} =
+  if newlen < x.len: shrink(x, newLen)
+  else: grow(x, newLen, default(T))
+
+template `[]`*[T](x: myseq[T]; i: Natural): T =
+  assert i < x.len
+  x.data[i]
+
+template `[]=`*[T](x: myseq[T]; i: Natural; y: T) =
+  assert i < x.len
+  x.data[i] = y
+
+proc createSeq*[T](elems: varargs[T]): myseq[T] =
+  result.cap = max(elems.len, 2)
+  result.len = elems.len
+  result.data = cast[type(result.data)](alloc0(result.cap * sizeof(T)))
+  inc allocCount
+  when supportsCopyMem(T):
+    copyMem(result.data, unsafeAddr(elems[0]), result.cap * sizeof(T))
+  else:
+    for i in 0..<result.len:
+      result.data[i] = elems[i]
+
+proc len*[T](x: myseq[T]): int {.inline.} = x.len
+
+proc main =
+  var s = createSeq(1, 2, 3, 4, 5, 6)
+  s.add 89
+  s.grow s.len + 2, 90
+  for i in 0 ..< s.len:
+    echo s[i]
+
+  var nested = createSeq(createSeq(1, 2, 3), createSeq(4, 5, 6, 7))
+  for i in 0 ..< nested.len:
+    for j in 0 ..< nested[i].len:
+      echo i, " ", j, " ", nested[i][j]
+
+main()
+echo "after ", allocCount, " ", deallocCount
+
+type
+  Node = ref object
+    name: char
+    sccId: int
+    kids: myseq[Node]
+    rc: int
+
+proc edge(a, b: Node) =
+  inc b.rc
+  a.kids.add b
+
+proc createNode(name: char): Node =
+  new result
+  result.name = name
+  result.kids = createSeq[Node]()
+
+proc use(x: Node) = discard
+
+proc buildComplexGraph: Node =
+  # see https://en.wikipedia.org/wiki/Strongly_connected_component for the
+  # graph:
+  let a = createNode('a')
+  let b = createNode('b')
+  let c = createNode('c')
+  let d = createNode('d')
+  let e = createNode('e')
+
+  a.edge c
+  c.edge b
+  c.edge e
+  b.edge a
+  d.edge c
+  e.edge d
+
+
+  let f = createNode('f')
+  b.edge f
+  e.edge f
+
+  let g = createNode('g')
+  let h = createNode('h')
+  let i = createNode('i')
+
+  f.edge g
+  f.edge i
+  g.edge h
+  h.edge i
+  i.edge g
+
+  let j = createNode('j')
+
+  h.edge j
+  i.edge j
+
+  let k = createNode('k')
+  let l = createNode('l')
+
+  f.edge k
+  k.edge l
+  l.edge k
+  k.edge j
+
+  let m = createNode('m')
+  let n = createNode('n')
+  let p = createNode('p')
+  let q = createNode('q')
+
+  m.edge n
+  n.edge p
+  n.edge q
+  q.edge p
+  p.edge m
+
+  q.edge k
+
+  d.edge m
+  e.edge n
+
+  result = a
+
+proc main2 =
+  let g = buildComplexGraph()
+
+main2()
+GC_fullCollect()
+echo "MEM ", getOccupiedMem()