#lang racket/base

(provide (struct-out annotated)
         annotate
         annotations
         strip-annotations
         strip-annotations-proc
         peel-annotations
         peel-annotations-proc)

(require racket/match)
(require (for-syntax racket/base))
(require "record.rkt")
(require racket/dict)
(require racket/set)

;; Syntax properties and syntax objects would be almost perfect for
;; representing annotations, plus position/source tracking as
;; lagniappe, but unfortunately they don't play nicely with data much
;; outside of ordinary S-expressions as found in Racket source.
;;
;; So we do our own thing, for now.
;;
;; See also https://gitlab.com/preserves/preserves/-/issues/16,
;; 'Consider Racket "correlated objects" for annotations
;; representation'.
;;
(struct annotated (annotations srcloc item) #:transparent
  #:methods gen:equal+hash
  [(define (equal-proc a b =?) (=? (annotated-item a) (annotated-item b)))
   (define (hash-proc a h) (h (annotated-item a)))
   (define (hash2-proc a h) (h (annotated-item a)))])

(define (annotate v . as)
  (match v
    [(annotated annotations srcloc item)
     (annotated (append as annotations) srcloc item)]
    [item
     (annotated as #f item)]))

(define (annotations v)
  (match v
    [(annotated annotations _ _) annotations]
    [_ '()]))

(define (strip-annotations-proc v #:depth [depth +inf.0])
  (let walk* ((v v) (depth depth))
    (define next-depth (- depth 1))
    (define (walk v) (walk* v next-depth))
    (if (zero? depth)
        v
        (match v
          [(annotated _ _ item)
           (match item
             [(record label fields) (record (walk* label depth) (map walk fields))]
             [(? list?) (map walk item)]
             [(? set?) (for/set [(i (in-set item))] (walk i))]
             [(? dict?) (for/hash [((k v) (in-dict item))] (values (walk* k depth) (walk v)))]
             [(? annotated?) (error 'strip-annotations "Improper annotation structure: ~v" v)]
             [_ item])]
          [_ v]))))

(define (peel-annotations-proc v)
  (strip-annotations-proc v #:depth 1))

(define-match-expander strip-annotations
  (syntax-rules () [(_ pat extra ...) (app (lambda (v) (strip-annotations-proc v extra ...)) pat)])
  (lambda (stx)
    (syntax-case stx ()
      [(_ args ...) #'(strip-annotations-proc args ...)]
      [_ #'strip-annotations-proc])))

(define-match-expander peel-annotations
  (syntax-rules () [(_ pat extra ...) (app (lambda (v) (peel-annotations-proc v extra ...)) pat)])
  (lambda (stx)
    (syntax-case stx ()
      [(_ args ...) #'(peel-annotations-proc args ...)]
      [_ #'peel-annotations-proc])))