;;; -*- Mode: Scheme; scheme48-package: statistical-profiling -*-

;;;; Scheme48 Statistical Profiler

;;; This code is written by Taylor Campbell and placed in the Public
;;; Domain.  All warranties are disclaimed.

(define-record-type* profile
  (make-profile
   )
  (
   (total-samples 0)
   (total-waiting 0)
   (call-graph (make-call-graph))
   (continuation #f)
   interrupt-handler
   ))

(define (with-profiling profile thunk)
  (dynamic-wind
    (lambda ()
      (start-profiling profile))
    (lambda ()
      (primitive-cwcc
       (lambda (continuation)
         ;; It would be nice if we could pull the modification of the
         ;; profiler's continuation out into the DYNAMIC-WIND thunks,
         ;; but we can't, because we want to exclude the DYNAMIC-WIND
         ;; continuation from the profiler's continuation.
         (set-profile-continuation! profile continuation)
         (receive vals (thunk)
           (set-profile-continuation! profile #f)
           (apply values vals)))))
    (lambda ()
      (stop-profiling profile))))

(define (profile-report profile)
  (let ((nodes '()))
    (walk-call-graph (lambda (node)
                       (set! nodes (cons node nodes)))
                     (profile-call-graph profile))
    (let* ((node-list (sort-list! nodes
                                  (lambda (x y)
                                    (> (profile-node-samples x)
                                       (profile-node-samples y)))))
           (node-count (length node-list))
           (node-vector (make-vector node-count)))
      (do ((i 0 (+ i 1))
           (l node-list (cdr l)))
          ((= i node-count))
        (vector-set! node-vector i (car l))
        (set-profile-node-index! (car l) i))
      node-vector)))

(define interrupt-handlers 1)           ;Magic constant!

(define (start-profiling profile)
  (let ((handlers (session-data-ref interrupt-handlers)))
    (set-profile-interrupt-handler!
     profile
     (vector-ref handlers (enum interrupt ALARM)))
    (vector-set! handlers
                 (enum interrupt ALARM)
                 (profiling-interrupt-handler profile))))

(define (stop-profiling profile)
  (let ((handler (profile-interrupt-handler profile)))
    (set-profile-interrupt-handler! profile #f)
    (vector-set! (session-data-ref interrupt-handlers)
                 (enum interrupt ALARM)
                 handler)))

(define (profiling-interrupt-handler profile)
  (let ((interrupt-handler (profile-interrupt-handler profile)))
    (lambda (template enabled-interrupts)
      (primitive-cwcc
       (lambda (continuation)
         ;; This continuation is a RETURN-FROM-INTERRUPT continuation,
         ;; so elide that by taking its parent.
         (let ((continuation (continuation-cont continuation)))
           (if (useful-continuation? continuation profile)
               (record-continuation continuation profile)))
         (interrupt-handler template enabled-interrupts))))))

(define (record-continuation continuation profile)
  (let* ((call-graph (profile-call-graph profile))
         (callee (find-sampled-node call-graph continuation)))
    (if callee
        (set-profile-total-samples!
         profile
         (+ 1 (profile-total-samples profile))))
    (let loop ((continuation continuation)
               (callee callee)
               (total-waiting (profile-total-waiting profile)))
      (let ((continuation (continuation-cont continuation)))
        (if (useful-continuation? continuation profile)
            (let ((caller
                   (find-waiting-node call-graph continuation)))
              (if (and caller callee)
                  (make-edge! caller callee))
              (loop continuation
                    caller
                    (if caller (+ total-waiting 1) total-waiting)))
            (set-profile-total-waiting! profile total-waiting))))))

(define (useful-continuation? continuation profile)
  (and continuation
       (not (eq? continuation (profile-continuation profile)))))

;;;; Call Graphs

(define (make-call-graph)
  (make-table template-id))

(define-record-type* profile-node
  (%make-profile-node
   template
   (samples)          ; number of direct samples
   (waiting)          ; number of times found waiting on the stack
   )
  (
   index
   (caller-edges (make-edges))
   (callee-edges (make-edges))
   ))

(define-record-discloser :profile-node
  (lambda (node)
    (list 'PROFILE-NODE
          (make-print-name
           (template-names (profile-node-template node)))
          (list 'SAMPLES: (profile-node-samples node))
          (list 'WAITING: (profile-node-waiting node)))))

(define (make-profile-node template samples waiting call-graph)
  (let ((node (%make-profile-node template samples waiting)))
    (set-call-graph-node! call-graph template node)
    node))

(define (call-graph-node call-graph template)
  (table-ref call-graph template))

(define (set-call-graph-node! call-graph template node)
  (table-set! call-graph template node))

(define (walk-call-graph procedure call-graph)
  (table-walk (lambda (template node)
                template
                (procedure node))
              call-graph))

(define (find-sampled-node call-graph continuation)
  (cond ((continuation-template continuation)
         => (lambda (template)
              (cond ((call-graph-node call-graph template)
                     => (lambda (node)
                          (set-profile-node-samples!
                           node
                           (+ 1 (profile-node-samples node)))
                          node))
                    (else
                     (make-profile-node template 1 0 call-graph)))))
        (else #f)))

(define (find-waiting-node call-graph continuation)
  (cond ((continuation-template continuation)
         => (lambda (template)
              (cond ((call-graph-node call-graph template)
                     => (lambda (node)
                          (set-profile-node-waiting!
                           node
                           (+ 1 (profile-node-waiting node)))
                          node))
                    (else
                     (make-profile-node template 0 1 call-graph)))))
        (else #f)))

;;;;; Call Graph Edges

(define (make-edges)
  (make-table template-id))

(define (make-edge! caller callee)
  (*make-edge! (profile-node-caller-edges callee)
               (profile-node-template caller))
  (*make-edge! (profile-node-callee-edges caller)
               (profile-node-template callee)))

(define (*make-edge! edges template)
  (table-set! edges
              template
              (+ 1 (or (table-ref edges template)
                       0))))

(define (profile-node-callers node profile)
  (let ((callers '()))
    (table-walk (let ((call-graph (profile-call-graph profile)))
                  (lambda (template count)
                    (set! callers
                          (cons (cons (call-graph-node call-graph
                                                       template)
                                      count)
                                callers))))
                (profile-node-caller-edges node))
    callers))

(define (profile-node-callees node profile)
  (let ((callees '()))
    (table-walk (let ((call-graph (profile-call-graph profile)))
                  (lambda (template count)
                    (set! callees
                          (cons (cons (call-graph-node call-graph
                                                       template)
                                      count)
                                callees))))
                (profile-node-callee-edges node))
    callees))