<commit_before>#lang racket

(require "alternative.rkt")
(require "programs.rkt")
(require "matcher.rkt")
(require "points.rkt")
(require "common.rkt")
(require "syntax.rkt")
(require "localize-error.rkt")

(provide infer-splitpoints (struct-out sp))

(define (infer-splitpoints alts)
  (debug "Finding splitpoints for:" alts #:from 'regime-changes #:depth 2)
  (let* ([options (map (curry option-on-expr alts)
		       (exprs-to-branch-on alts))]
	 [best-option (argmin (compose errors-score option-errors) options)]
	 [splitpoints (option-splitpoints best-option)]
	 [altns (used-alts splitpoints alts)]
	 [splitpoints* (coerce-indices splitpoints)])
    (debug "Found splitpoints:" splitpoints* ", with alts" altns)
    (list splitpoints* altns)))

(struct option (splitpoints errors) #:transparent
	#:methods gen:custom-write
        [(define (write-proc opt port mode)
           (display "#<option " port)
           (write (option-splitpoints opt) port)
           (display ">" port))])

(define (exprs-to-branch-on alts)
  (cons (critical-subexpression (*start-prog*))
	(program-variables (alt-program (car alts)))))

(define (critical-subexpression prog)
  (define (loc-children loc subexpr)
    (map (compose (curry append loc)
		  list)
	 (range 1 (length subexpr))))
  (define (all-equal? items)
    (if (< (length items) 2) #t
	(and (equal? (car items) (cadr items)) (all-equal? (cdr items)))))
  (define (critical-child expr)
    (let ([var-locs
	   (let get-vars ([subexpr expr]
			  [cur-loc '()])
	     (cond [(list? subexpr)
		    (append-map get-vars (cdr subexpr)
				(loc-children cur-loc subexpr))]
		   [(constant? subexpr)
		    '()]
		   [(variable? subexpr)
		    (list (cons subexpr cur-loc))]))])
      (if (all-equal? (map car var-locs))
	  (caar var-locs)
	  (let get-subexpr ([subexpr expr] [vlocs var-locs])
	    (cond [(all-equal? (map cadr vlocs))
		   (get-subexpr (if (= 1 (cadar vlocs)) (cadr subexpr) (caddr subexpr))
				(for/list ([vloc vlocs])
				  (cons (car vloc) (cddr vloc))))]
		  [#t subexpr])))))
  (let* ([loc (car (localize-error prog))])
    (critical-child (location-get loc prog))))

(define basic-point-search (curry binary-search (λ (p1 p2)
						  (if (for/and ([val1 p1] [val2 p2])
							(> *epsilon-fraction* (abs (- val1 val2))))
						      p1
						      (for/list ([val1 p1] [val2 p2])
							(/ (+ val1 val2) 2))))))

(define (used-alts splitpoints all-alts)
  (let ([used-indices (remove-duplicates (map sp-cidx splitpoints))])
    (map (curry list-ref all-alts) used-indices)))

;; Takes a list of splitpoints, `splitpoints`, whose indices originally referred to some list of alts `alts`,
;; and changes their indices so that they are consecutive starting from zero, but all indicies that
;; previously matched still match.
(define (coerce-indices splitpoints)
  (let* ([used-indices (remove-duplicates (map sp-cidx splitpoints))]
	 [mappings (map cons used-indices (range (length used-indices)))])
    (map (λ (splitpoint)
	   (sp (cdr (assoc (sp-cidx splitpoint) mappings))
	       (sp-bexpr splitpoint)
	       (sp-point splitpoint)))
	 splitpoints)))

(define (option-on-expr alts expr)
  (match-let* ([vars (program-variables (*start-prog*))]
	       [`(,pts ,exs) (sort-context-on-expr (*pcontext*) expr vars)])
    (let* ([err-lsts (parameterize ([*pcontext* (mk-pcontext pts exs)])
		       (map alt-errors alts))]
	   [bit-err-lsts (map (curry map ulps->bits) err-lsts)]
	   [split-indices (err-lsts->split-indices bit-err-lsts)]
	   [split-points (sindices->spoints pts expr alts split-indices)])
      (option split-points (pick-errors split-points pts err-lsts vars)))))

(define (error-at prog point exact)
  (car (errors prog (mk-pcontext (vector point) (vector exact)))))

;; Accepts a list of sindices in one indexed form and returns the
;; proper splitpoints in floath form.
(define (sindices->spoints points expr alts sindices)
  (define (eval-on-pt pt)
    ((eval-prog `(λ ,(program-variables (alt-program (car alts))) ,expr) mode:fl) pt))
  (define (sidx->spoint sidx next-sidx)
    (let* ([alt1 (list-ref alts (si-cidx sidx))]
	   [alt2 (list-ref alts (si-cidx next-sidx))]
	   [pred (λ (p)
		   (let* ([exact ((eval-prog (*start-prog*) mode:bf) p)]
			  [e1 (error-at (alt-program alt1) p exact)]
			  [e2 (error-at (alt-program alt2) p exact)])
		     (< e1 e2)))])
      (sp (si-cidx sidx) expr (eval-on-pt (basic-point-search pred (list-ref points (si-pidx sidx)) (list-ref points (sub1 (si-pidx sidx))))))))
  (append (map sidx->spoint
	       (take sindices (sub1 (length sindices)))
	       (drop sindices 1))
	  (list (let ([last-sidx (list-ref sindices (sub1 (length sindices)))])
		  (sp (si-cidx last-sidx)
		      expr
		      +inf.0)))))

(define (point-with-dim index point val)
  (map (λ (pval pindex) (if (= pindex index) val pval))
       point
       (range (length point))))

(define (pick-errors splitpoints pts err-lsts variables)
  (reverse
   (first-value
    (for/fold ([acc '()] [rest-splits splitpoints])
	([pt (in-list pts)]
	 [errs (flip-lists err-lsts)])
      (if (<= ((eval-prog `(λ ,variables ,(sp-bexpr (car rest-splits))) mode:fl) pt)
	      (sp-point (car rest-splits)))
	  (values (cons (list-ref errs (sp-cidx (car rest-splits)))
			acc)
		  rest-splits)
	  (values acc (cdr rest-splits)))))))<commit_msg>Don't try to binary search on expression<commit_after>#lang racket

(require "alternative.rkt")
(require "programs.rkt")
(require "matcher.rkt")
(require "points.rkt")
(require "common.rkt")
(require "syntax.rkt")
(require "localize-error.rkt")

(provide infer-splitpoints (struct-out sp))

(define (infer-splitpoints alts)
  (debug "Finding splitpoints for:" alts #:from 'regime-changes #:depth 2)
  (let* ([options (map (curry option-on-expr alts)
		       (exprs-to-branch-on alts))]
	 [best-option (argmin (compose errors-score option-errors) options)]
	 [splitpoints (option-splitpoints best-option)]
	 [altns (used-alts splitpoints alts)]
	 [splitpoints* (coerce-indices splitpoints)])
    (debug "Found splitpoints:" splitpoints* ", with alts" altns)
    (list splitpoints* altns)))

(struct option (splitpoints errors) #:transparent
	#:methods gen:custom-write
        [(define (write-proc opt port mode)
           (display "#<option " port)
           (write (option-splitpoints opt) port)
           (display ">" port))])

(define (exprs-to-branch-on alts)
  (cons (critical-subexpression (*start-prog*))
	(program-variables (alt-program (car alts)))))

(define (critical-subexpression prog)
  (define (loc-children loc subexpr)
    (map (compose (curry append loc)
		  list)
	 (range 1 (length subexpr))))
  (define (all-equal? items)
    (if (< (length items) 2) #t
	(and (equal? (car items) (cadr items)) (all-equal? (cdr items)))))
  (define (critical-child expr)
    (let ([var-locs
	   (let get-vars ([subexpr expr]
			  [cur-loc '()])
	     (cond [(list? subexpr)
		    (append-map get-vars (cdr subexpr)
				(loc-children cur-loc subexpr))]
		   [(constant? subexpr)
		    '()]
		   [(variable? subexpr)
		    (list (cons subexpr cur-loc))]))])
      (if (all-equal? (map car var-locs))
	  (caar var-locs)
	  (let get-subexpr ([subexpr expr] [vlocs var-locs])
	    (cond [(all-equal? (map cadr vlocs))
		   (get-subexpr (if (= 1 (cadar vlocs)) (cadr subexpr) (caddr subexpr))
				(for/list ([vloc vlocs])
				  (cons (car vloc) (cddr vloc))))]
		  [#t subexpr])))))
  (let* ([loc (car (localize-error prog))])
    (critical-child (location-get loc prog))))

(define basic-point-search (curry binary-search (λ (p1 p2)
						  (if (for/and ([val1 p1] [val2 p2])
							(> *epsilon-fraction* (abs (- val1 val2))))
						      p1
						      (for/list ([val1 p1] [val2 p2])
							(/ (+ val1 val2) 2))))))

(define (used-alts splitpoints all-alts)
  (let ([used-indices (remove-duplicates (map sp-cidx splitpoints))])
    (map (curry list-ref all-alts) used-indices)))

;; Takes a list of splitpoints, `splitpoints`, whose indices originally referred to some list of alts `alts`,
;; and changes their indices so that they are consecutive starting from zero, but all indicies that
;; previously matched still match.
(define (coerce-indices splitpoints)
  (let* ([used-indices (remove-duplicates (map sp-cidx splitpoints))]
	 [mappings (map cons used-indices (range (length used-indices)))])
    (map (λ (splitpoint)
	   (sp (cdr (assoc (sp-cidx splitpoint) mappings))
	       (sp-bexpr splitpoint)
	       (sp-point splitpoint)))
	 splitpoints)))

(define (option-on-expr alts expr)
  (match-let* ([vars (program-variables (*start-prog*))]
	       [`(,pts ,exs) (sort-context-on-expr (*pcontext*) expr vars)])
    (let* ([err-lsts (parameterize ([*pcontext* (mk-pcontext pts exs)])
		       (map alt-errors alts))]
	   [bit-err-lsts (map (curry map ulps->bits) err-lsts)]
	   [split-indices (err-lsts->split-indices bit-err-lsts)]
	   [split-points (sindices->spoints pts expr alts split-indices)])
      (option split-points (pick-errors split-points pts err-lsts vars)))))

(define (error-at prog point exact)
  (car (errors prog (mk-pcontext (vector point) (vector exact)))))

;; Accepts a list of sindices in one indexed form and returns the
;; proper splitpoints in floath form.
(define (sindices->spoints points expr alts sindices)
  (define (eval-on-pt pt)
    ((eval-prog `(λ ,(program-variables (alt-program (car alts))) ,expr) mode:fl) pt))
  (define (sidx->spoint sidx next-sidx)
    ;; Todo: Do something smart for expressions.
    (if (list? expr) (sp (si-cidx sidx) expr (eval-on-pt (list-ref points (si-pidx sidx))))
	(let* ([alt1 (list-ref alts (si-cidx sidx))]
	       [alt2 (list-ref alts (si-cidx next-sidx))]
	       [pred (λ (p)
		       (let* ([exact ((eval-prog (*start-prog*) mode:bf) p)]
			      [e1 (error-at (alt-program alt1) p exact)]
			      [e2 (error-at (alt-program alt2) p exact)])
			 (< e1 e2)))])
	  (sp (si-cidx sidx) expr (eval-on-pt (basic-point-search
					       pred (list-ref points (si-pidx sidx))
					       (list-ref points (sub1 (si-pidx sidx)))))))))
  (append (map sidx->spoint
	       (take sindices (sub1 (length sindices)))
	       (drop sindices 1))
	  (list (let ([last-sidx (list-ref sindices (sub1 (length sindices)))])
		  (sp (si-cidx last-sidx)
		      expr
		      +inf.0)))))

(define (point-with-dim index point val)
  (map (λ (pval pindex) (if (= pindex index) val pval))
       point
       (range (length point))))

(define (pick-errors splitpoints pts err-lsts variables)
  (reverse
   (first-value
    (for/fold ([acc '()] [rest-splits splitpoints])
	([pt (in-list pts)]
	 [errs (flip-lists err-lsts)])
      (if (<= ((eval-prog `(λ ,variables ,(sp-bexpr (car rest-splits))) mode:fl) pt)
	      (sp-point (car rest-splits)))
	  (values (cons (list-ref errs (sp-cidx (car rest-splits)))
			acc)
		  rest-splits)
	  (values acc (cdr rest-splits)))))))
{'fraction_non_alphanumeric': 0.16743372606774667, 'fraction_numerical': 0.006075110456553756, 'mean_word_length': 3.3792825473599355, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 19, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1325425', 'n_tokens_mistral': 4173, 'n_tokens_neox': 3818, 'n_words': 1113}