<commit_before>#lang racket

(require "header.rkt" "arrayforth.rkt" "arrayforth-print.rkt" "arrayforth-def.rkt")
(require "../../forth-interpreter/machine/cegis.rkt" 
         "../../forth-interpreter/machine/state.rkt" 
         "../../forth-interpreter/machine/programs.rkt" 
         "../../forth-interpreter/machine/track-constant.rkt"
         "../../forth-interpreter/ArrayForth/arrayforth.rkt")
;; (require (planet "machine/cegis.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "machine/state.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "machine/programs.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "machine/track-constant.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "ArrayForth/arrayforth.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))

(provide superoptimize renameindex program-sizes)

(define (out-space out cnstr #:mem [mem #t])
  (define result
    (if (= out 0)
	(constraint s t r)
	(constraint (data out) s t r)))
  (struct-copy progstate result 
               [memory (and mem (restrict-mem cnstr))] 
               [a (restrict-a cnstr)]
               [b (restrict-b cnstr)]
               [r #t]
               [return (if (restrict-r cnstr) 2 2)]
               ))

(define (in-constraint data inst)
  (cond
   [(equal? data #f) 
    (default-state)]
   [(or (equal? inst "b!") (equal? inst "a!")) 
    (default-state (t (port-number data)))]
   [(or (equal? inst "!b") (equal? inst "@b")) 
    (default-state (b (port-number data)))]
   [(or (equal? inst "!") (equal? inst "@")) 
    (default-state (a (port-number data)))]
   [else
    (default-state)]))

(define index-map #f)
(define mem-size #f)
(define bit #f)
(define simple #f)

(define name #f)
(define w #f)
(define h #f)
(define id #f)
(define before 0)
(define after 0)
(define sliding #t)
(define dir #f)

(define (superoptimize ast my-name my-w my-h sliding-windows my-dir #:id [id 0])
  (set! name my-name)
  (set! w my-w)
  (set! h my-h)
  (set! id id)
  (set! before 0)
  (set! after 0)
  (set! sliding sliding-windows)
  (set! dir my-dir)
  (system (format "rm ~a/~a.stat" dir name))
  (system (format "rm ~a/~a-work.rkt" dir name))
  (system (format "rm ~a/~a-work.aforth" dir name))
  (let ([result (superoptimize-inner ast)])
    ;; (with-output-to-file #:exists 'append 
    ;;                      (format "~a/~a.stat" dir name)
    ;;                      (lambda () 
    ;;                        (pretty-display (format "~a ~a (TOTAL)" before after))))
    result))

;; optimize per-core program
(define (superoptimize-inner ast)
  (define (collect-code x big-block limit)
    (define entry (linklist-entry x))
    (cond
     [(block? entry)
      (define body (block-body entry))
      (define seq (if (string? body) (string-split body) body))
      
      (if (or (= (length (block-body big-block)) 0)
              (<= (+ (length (block-body big-block)) (length seq)) limit))
          (collect-code (linklist-next x) (merge-block big-block entry) limit)
          (values x big-block))]
     
     [else (values x big-block)]))
  
  (cond
   [(equal? ast #f) #f]

   [(block? ast)
    (pretty-display "superoptimize: block")
    ;; (raise "superoptimizer: optimize via linklist not block!")
    (define out (block-out ast))
    (define body-list (if (string? (block-body ast)) 
                          (string-split (block-body ast))
                          (block-body ast)))
                     
    (define result (optimize (string-join body-list)
                             #:f18a #f
                             #:num-bits max-bit #:name name
                             #:constraint (out-space out (block-cnstr ast))
                             #:start-state (if (empty? body-list) 
                                               (default-state)
                                               (in-constraint (block-incnstr ast) 
                                                              (car body-list)))
                             #:mem mem-size #:start mem-size #:bin-search `length))

    (pretty-display `(result ,result))
    (define opt (if (equal? result 'timeout)
                    ast
                    (new-block result (block-in ast) out (block-cnstr ast) (block-org ast))))

    (define renamed (renameindex opt mem-size index-map))
    (pretty-display "CASE 3")
    (aforth-syntax-print renamed 1 1)
    
    (with-output-to-file #:exists 'append 
      (format "~a/~a-work.rkt" dir name)
      (lambda () 
        (pretty-display ";; original")
        (aforth-struct-print ast)
        (pretty-display ";; optimized")
        (aforth-struct-print opt)
        (pretty-display ";; result")
        (aforth-struct-print renamed)))

    renamed
    ]

   [(linklist? ast)
    (pretty-display (format "superoptimize: linklist ~a" (linklist-entry ast)))
    (cond
     [(and (linklist-entry ast) (block? (linklist-entry ast)))
      ;; optimize code in group of blocks whose length <= len
      ;; if timeout => reduce len
      ;; if getting result => return the optimized block
      (define (optimize-loop len)
        (define-values (next block-noopt) 
          (collect-code ast (new-block (list) 0 0 (restrict #t #f #f #f) (list)) len))
        
        (define out (block-out block-noopt))
        (define body (block-body block-noopt))
        (pretty-display `(optimize-loop ,len ,body ,sliding))
        (define cnstr (block-cnstr block-noopt))
        (define start-state (if (empty? body)
                                (default-state)
                                (in-constraint 
                                 (block-incnstr block-noopt)
                                 (car body))))
        (define result (optimize (string-join body)
				 #:f18a #f
				 #:num-bits max-bit #:name name
				 ;; Not constraining mem when len > block-limit because
				 ;; we optimize the entire function.
				 #:constraint (out-space out cnstr #:mem (<= len block-limit))
                                 #:start-state (constrain-stack start-state simple)
				 #:mem mem-size #:start mem-size #:bin-search `length))

        (cond
         [(and sliding (equal? result 'timeout))
          (let* ([last-block (linklist-entry (linklist-prev next))]
                 [last-body (block-body last-block)]
                 [last-list (if (string? last-body) (string-split last-body) last-body)]
                 [new-len (if (> len block-limit)
			      block-limit
			      (- (length body) (max 1 (length last-list))))])
            (if (> new-len 0)
                (optimize-loop new-len)
                (begin
		  (set! simple #f)
		  (values next block-noopt))))]

         [(equal? result 'timeout)
	  (set! simple #f)
          (values next block-noopt)]

         [else
	  (set! simple #f)
          (set-block-body! block-noopt result)
          (values next block-noopt)]))

      ;; If simple function, set limit to more than max.
      (define-values (next opt) (optimize-loop (if #f 1000 block-limit)))
      (pretty-display "OPTMIZE: DONE")
      (define renamed (renameindex opt mem-size index-map))
      (pretty-display "RENAME: DONE")
    <commit_msg>pass id to superoptimizer<commit_after>#lang racket

(require "header.rkt" "arrayforth.rkt" "arrayforth-print.rkt" "arrayforth-def.rkt")
(require "../../forth-interpreter/machine/cegis.rkt" 
         "../../forth-interpreter/machine/state.rkt" 
         "../../forth-interpreter/machine/programs.rkt" 
         "../../forth-interpreter/machine/track-constant.rkt"
         "../../forth-interpreter/ArrayForth/arrayforth.rkt")
;; (require (planet "machine/cegis.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "machine/state.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "machine/programs.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "machine/track-constant.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))
;; (require (planet "ArrayForth/arrayforth.rkt" ("mangpo" "aforth-optimizer.plt" 1 0)))

(provide superoptimize renameindex program-sizes)

(define (out-space out cnstr #:mem [mem #t])
  (define result
    (if (= out 0)
	(constraint s t r)
	(constraint (data out) s t r)))
  (struct-copy progstate result 
               [memory (and mem (restrict-mem cnstr))] 
               [a (restrict-a cnstr)]
               [b (restrict-b cnstr)]
               [r #t]
               [return (if (restrict-r cnstr) 2 2)]
               ))

(define (in-constraint data inst)
  (cond
   [(equal? data #f) 
    (default-state)]
   [(or (equal? inst "b!") (equal? inst "a!")) 
    (default-state (t (port-number data)))]
   [(or (equal? inst "!b") (equal? inst "@b")) 
    (default-state (b (port-number data)))]
   [(or (equal? inst "!") (equal? inst "@")) 
    (default-state (a (port-number data)))]
   [else
    (default-state)]))

(define index-map #f)
(define mem-size #f)
(define bit #f)
(define simple #f)

(define name #f)
(define w #f)
(define h #f)
(define id #f)
(define before 0)
(define after 0)
(define sliding #t)
(define dir #f)

(define (superoptimize ast my-name my-w my-h sliding-windows my-dir #:id [my-id 0])
  (set! name my-name)
  (set! w my-w)
  (set! h my-h)
  (set! id my-id)
  (set! before 0)
  (set! after 0)
  (set! sliding sliding-windows)
  (set! dir my-dir)
  (system (format "rm ~a/~a.stat" dir name))
  (system (format "rm ~a/~a-work.rkt" dir name))
  (system (format "rm ~a/~a-work.aforth" dir name))
  (let ([result (superoptimize-inner ast)])
    ;; (with-output-to-file #:exists 'append 
    ;;                      (format "~a/~a.stat" dir name)
    ;;                      (lambda () 
    ;;                        (pretty-display (format "~a ~a (TOTAL)" before after))))
    result))

;; optimize per-core program
(define (superoptimize-inner ast)
  (define (collect-code x big-block limit)
    (define entry (linklist-entry x))
    (cond
     [(block? entry)
      (define body (block-body entry))
      (define seq (if (string? body) (string-split body) body))
      
      (if (or (= (length (block-body big-block)) 0)
              (<= (+ (length (block-body big-block)) (length seq)) limit))
          (collect-code (linklist-next x) (merge-block big-block entry) limit)
          (values x big-block))]
     
     [else (values x big-block)]))
  
  (cond
   [(equal? ast #f) #f]

   [(block? ast)
    (pretty-display "superoptimize: block")
    ;; (raise "superoptimizer: optimize via linklist not block!")
    (define out (block-out ast))
    (define body-list (if (string? (block-body ast)) 
                          (string-split (block-body ast))
                          (block-body ast)))
                     
    (define result (optimize (string-join body-list)
                             #:f18a #f
                             #:num-bits max-bit #:name name
                             #:constraint (out-space out (block-cnstr ast))
                             #:start-state (if (empty? body-list) 
                                               (default-state)
                                               (in-constraint (block-incnstr ast) 
                                                              (car body-list)))
                             #:mem mem-size #:start mem-size
                             #:bin-search `length #:id id))

    (pretty-display `(result ,result))
    (define opt (if (equal? result 'timeout)
                    ast
                    (new-block result (block-in ast) out (block-cnstr ast) (block-org ast))))

    (define renamed (renameindex opt mem-size index-map))
    (pretty-display "CASE 3")
    (aforth-syntax-print renamed 1 1)
    
    (with-output-to-file #:exists 'append 
      (format "~a/~a-work.rkt" dir name)
      (lambda () 
        (pretty-display ";; original")
        (aforth-struct-print ast)
        (pretty-display ";; optimized")
        (aforth-struct-print opt)
        (pretty-display ";; result")
        (aforth-struct-print renamed)))

    renamed
    ]

   [(linklist? ast)
    (pretty-display (format "superoptimize: linklist ~a" (linklist-entry ast)))
    (cond
     [(and (linklist-entry ast) (block? (linklist-entry ast)))
      ;; optimize code in group of blocks whose length <= len
      ;; if timeout => reduce len
      ;; if getting result => return the optimized block
      (define (optimize-loop len)
        (define-values (next block-noopt) 
          (collect-code ast (new-block (list) 0 0 (restrict #t #f #f #f) (list)) len))
        
        (define out (block-out block-noopt))
        (define body (block-body block-noopt))
        (pretty-display `(optimize-loop ,len ,body ,sliding ,id))
        (define cnstr (block-cnstr block-noopt))
        (define start-state (if (empty? body)
                                (default-state)
                                (in-constraint 
                                 (block-incnstr block-noopt)
                                 (car body))))
        (define result (optimize (string-join body)
				 #:f18a #f
				 #:num-bits max-bit #:name name
				 ;; Not constraining mem when len > block-limit because
				 ;; we optimize the entire function.
				 #:constraint (out-space out cnstr #:mem (<= len block-limit))
                                 #:start-state (constrain-stack start-state simple)
				 #:mem mem-size #:start mem-size
                                 #:bin-search `length
                                 #:id id))

        (cond
         [(and sliding (equal? result 'timeout))
          (let* ([last-block (linklist-entry (linklist-prev next))]
                 [last-body (block-body last-block)]
                 [last-list (if (string? last-body) (string-split last-body) last-body)]
                 [new-len (if (> len block-limit)
			      block-limit
			      (- (length body) (max 1 (length last-list))))])
            (if (> new-len 0)
                (optimize-loop new-len)
                (begin
		  (set! simple #f)
		  (values next block-noopt))))]

         [(equal? result 'timeout)
	  (set! simple #f)
          (values next block-noopt)]

         [else
	  (set! simple #f)
          (set-block-body! block-noopt result)
          (values next block-noopt)]))

      ;; If simple function, set limit to more than max.
      (define-values (next opt) (optimize-loop (if #f 1000 block-limit)))
      (pretty-display "OPTMIZE: DONE")
      (define renamed (renameindex opt mem-size index-map))
      (pretty-display "RENAME: DONE")
    
{'fraction_non_alphanumeric': 0.15010687443977108, 'fraction_numerical': 0.004688685099634558, 'mean_word_length': 1.9384116693679092, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 20, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6172344', 'n_tokens_mistral': 4861, 'n_tokens_neox': 4489, 'n_words': 1278}