<commit_before>#lang racket/base
(provide
  =/=
  =/=*
  absento
  numbero
  symbolo
  run-da
  run-da-depth
  run*-da
  run*-da-depth
  run-da-dls
  run*-da-dls
  )

(require
  "dict.rkt"
  "maybe.rkt"
  "microkanren.rkt"
  "minikanren.rkt"
  "record.rkt"
  "set.rkt"
  "sugar.rkt"
  racket/function
  (except-in racket/match ==)
  racket/set
  )

(module+ test
  (require
    rackunit
    ))

(define diseq-cx->var caar)
(define absent-cx->var cadr)

(record constraints diseqs absents type)
(define constraints-empty-v (constraints set-empty set-empty #f))
(define constraints-empty-p (constraints '() '() '()))
(define heq-empty (hasheq))

(record da-constraints pending var=>cxs)
(define da-constraints-empty (da-constraints constraints-empty-p heq-empty))
(def (da-constraints-diseqs-pending-add
       (da-constraints (constraints ds as ts) vs) cx)
  (da-constraints (constraints (cons cx ds) as ts) vs))
(def (da-constraints-absents-pending-add
       (da-constraints (constraints ds as ts) vs) cx)
  (da-constraints (constraints ds (cons cx as) ts) vs))
(def (da-constraints-types-pending-add
       (da-constraints (constraints ds as ts) vs) cx)
  (da-constraints (constraints ds as (cons cx ts)) vs))
(def (constraints-pending-add (constraints ds as ts)
                              (constraints ds0 as0 type) vr)
  ; TODO: use vr to rebuild cxs once they're specialized in a later diff
  (constraints (append (set->list ds0) ds)
               (append (set->list as0) as)
               (if type (cons (list type vr) ts) ts)))
(def (da-constraints-pending-clear (da-constraints _ vs))
  (da-constraints constraints-empty-p vs))

(def (da-constraints-var=>cxs-set (da-constraints p _) vc)
  (da-constraints p vc))
(define (da-constraints-new-bindings dats vrs)
  (if (null? vrs) dats
    (lets (da-constraints pending var=>cxs) = dats
          (values pending var=>cxs) =
          (forf pending = pending var=>cxs = var=>cxs
                vr <- vrs
                (match (hash-get var=>cxs vr)
                  ((nothing) (values pending var=>cxs))
                  ((just cxs) (values (constraints-pending-add pending cxs vr)
                                      (hash-remove var=>cxs vr)))))
          (da-constraints pending var=>cxs))))

(define (da-state-constraints-diseqs-pending-add st cx)
  (muk-state-constraints-set
    st (da-constraints-diseqs-pending-add (muk-state-constraints st) cx)))
(define (da-state-constraints-absents-pending-add st cx)
  (muk-state-constraints-set
    st (da-constraints-absents-pending-add (muk-state-constraints st) cx)))
(define (da-state-constraints-types-pending-add st cx)
  (muk-state-constraints-set
    st (da-constraints-types-pending-add (muk-state-constraints st) cx)))

(def (constraints-diseqs-add (constraints ds as ts) cx)
  (constraints (set-add ds cx) as ts))
(def (constraints-absents-add (constraints ds as ts) cx)
  (constraints ds (set-add as cx) ts))
(def (constraints-type-set (constraints ds as _) type)
  (constraints ds as type))

(define da-state-empty (muk-state-empty/constraints da-constraints-empty))
(def (da-add-constraint st name args)
  (match name
    ('=/=* (da-state-constraints-diseqs-pending-add st args))
    ('absento (da-state-constraints-absents-pending-add st args))
    ('type (da-state-constraints-types-pending-add st args))
    (_ (error (format "unsupported constraint: ~a ~a" name args)))))
(def (da-constrain st)
  (values st vr-new) = (muk-sub-new-bindings st)
  dats = (da-constraints-new-bindings (muk-state-constraints st) vr-new)
  pending = (da-constraints-pending dats)
  dats =
  (if (null? pending) dats
    (lets
      (constraints ds as ts) = pending

      (values st var=>cxs) =
      (da-constrain-types st (da-constraints-var=>cxs dats) ts)

      (and var=>cxs
           (forf dats = (da-constraints-var=>cxs-set
                          (da-constraints-pending-clear dats) var=>cxs)
                 (list cadd constrain cx->var pending) <-
                 (list (list constraints-diseqs-add
                             da-constrain-diseq
                             diseq-cx->var
                             ds)
                       (list constraints-absents-add
                             da-constrain-absent
                             absent-cx->var
                             as))
                 #:break (not dats)
                 new = (let loop ((new '()) (pending pending))
                         (match pending
                           ('() new)
                           ((cons cx pending)
                            (lets cxs = (constrain st cx)
                                  (and cxs (loop (append cxs new) pending))))))
                 (and new
                      (da-constraints-var=>cxs-set
                        dats
                        (forf var=>cxs = (da-constraints-var=>cxs dats)
                              cx <- new
                              (hash-update
                                var=>cxs
                                (cx->var cx)
                                (lambda (cxs) (cadd cxs cx))
                                constraints-empty-v))))))))
  (or (and dats (list (muk-state-constraints-set st dats))) '()))

(define (da-constrain-diseq st or-diseqs)
  (def (muk-var< (muk-var n0) (muk-var n1)) (symbol<? n0 n1))
  (def (total< e0 e1)
    (or (not (muk-var? e1)) (and (muk-var? e0) (muk-var< e0 e1))))
  (def (cons< (cons k0 v0) (cons k1 v1)) (muk-var< k0 k1))
  (let ((st (forf st = st
                  (cons e0 e1) <- or-diseqs
                  #:break (not st)
                  (muk-unify st e0 e1))))
    (if st
      (lets (values st vr-new) = (muk-sub-new-bindings st)
            or-diseqs = (forl vr <- vr-new
                              (values _ val) = (muk-sub-get st vr)
                              (apply cons (sort (list vr val) total<)))
            or-diseqs = (sort or-diseqs cons<)
            (and (pair? or-diseqs) (list or-diseqs)))
      '())))
(define (=/=* or-diseqs) (muk-constraint '=/=* or-diseqs))
(define (=/= e0 e1) (=/=* `((,e0 . ,e1))))

(def (da-constrain-absent st (list ground tm))
  (values st tm) = (muk-walk st tm)
  (match tm
    ((? muk-var?) (list (list ground tm)))
    ((cons h0 t0)
     (let ((next0 (da-constrain-absent st (list ground h0))))
       (and next0
            (let ((next1 (da-constrain-absent st (list ground t0))))
              (and next1 (append next0 next1))))))
    (_ (if (eqv? ground tm) #f '()))))
(define (absento ground tm)
  (if (pair? ground)
    (error (format "absento only supports absence of ground terms: ~a ~a"
                   ground tm))
    (muk-constraint 'absento (list ground tm))))

(define (da-constrain-types st var=>cxs ts)
  (forf st = st var=>cxs = var=>cxs
        (list tag tm) <- ts
        #:break (not var=>cxs)
        (values st tm) = (muk-walk st tm)
        (values st
          (if (muk-var? tm)
            (lets
              cxs = (hash-ref var=>cxs tm constraints-empty-v)
              tag0 = (constraints-type cxs)
              (if tag0
                (and (eq? tag0 tag) var=>cxs)
                (hash-set var=>cxs tm (constraints-type-set cxs tag))))
            (and (or (and (symbol? tm) (eq? tag 'sym))
                     (and (number? tm) (eq? tag 'num))) var=>cxs)))))
(define (typeo tag tm)
  (if (set-member? (set 'num 'sym) tag) (muk-constraint 'type (list tag tm))
    (error (format "invalid type tag: ~a ~a" tag tm))))
(define (symbolo tm) (typeo 'sym tm))
(define (numbero tm) (typeo 'num tm))<commit_msg>Process absento constraints more efficiently<commit_after>#lang racket/base
(provide
  =/=
  =/=*
  absento
  numbero
  symbolo
  run-da
  run-da-depth
  run*-da
  run*-da-depth
  run-da-dls
  run*-da-dls
  )

(require
  "dict.rkt"
  "maybe.rkt"
  "microkanren.rkt"
  "minikanren.rkt"
  "record.rkt"
  "set.rkt"
  "sugar.rkt"
  racket/function
  (except-in racket/match ==)
  racket/set
  )

(module+ test
  (require
    rackunit
    ))

(define diseq-cx->var caar)

(record constraints diseqs absents type)
(define constraints-empty-v (constraints set-empty set-empty #f))
(define constraints-empty-p (constraints '() '() '()))
(define heq-empty (hasheq))

(record da-constraints pending var=>cxs)
(define da-constraints-empty (da-constraints constraints-empty-p heq-empty))
(def (da-constraints-diseqs-pending-add
       (da-constraints (constraints ds as ts) vs) cx)
  (da-constraints (constraints (cons cx ds) as ts) vs))
(def (da-constraints-absents-pending-add
       (da-constraints (constraints ds as ts) vs) cx)
  (da-constraints (constraints ds (cons cx as) ts) vs))
(def (da-constraints-types-pending-add
       (da-constraints (constraints ds as ts) vs) cx)
  (da-constraints (constraints ds as (cons cx ts)) vs))
(def (constraints-pending-add (constraints ds as ts)
                              (constraints ds0 as0 type) vr)
  ; TODO: use vr to rebuild cxs once they're specialized in a later diff
  (constraints (append (set->list ds0) ds)
               (append (set-map as0 (lambda (ground) (list ground vr))) as)
               (if type (cons (list type vr) ts) ts)))
(def (da-constraints-pending-clear (da-constraints _ vs))
  (da-constraints constraints-empty-p vs))

(def (da-constraints-var=>cxs-set (da-constraints p _) vc)
  (da-constraints p vc))
(define (da-constraints-new-bindings dats vrs)
  (if (null? vrs) dats
    (lets (da-constraints pending var=>cxs) = dats
          (values pending var=>cxs) =
          (forf pending = pending var=>cxs = var=>cxs
                vr <- vrs
                (match (hash-get var=>cxs vr)
                  ((nothing) (values pending var=>cxs))
                  ((just cxs) (values (constraints-pending-add pending cxs vr)
                                      (hash-remove var=>cxs vr)))))
          (da-constraints pending var=>cxs))))

(define (da-state-constraints-diseqs-pending-add st cx)
  (muk-state-constraints-set
    st (da-constraints-diseqs-pending-add (muk-state-constraints st) cx)))
(define (da-state-constraints-absents-pending-add st cx)
  (muk-state-constraints-set
    st (da-constraints-absents-pending-add (muk-state-constraints st) cx)))
(define (da-state-constraints-types-pending-add st cx)
  (muk-state-constraints-set
    st (da-constraints-types-pending-add (muk-state-constraints st) cx)))

(def (constraints-diseqs-add (constraints ds as ts) cx)
  (constraints (set-add ds cx) as ts))
(def (constraints-absents-add (constraints ds as ts) ground)
  (constraints ds (set-add as ground) ts))
(def (constraints-type-set (constraints ds as _) type)
  (constraints ds as type))

(define da-state-empty (muk-state-empty/constraints da-constraints-empty))
(def (da-add-constraint st name args)
  (match name
    ('=/=* (da-state-constraints-diseqs-pending-add st args))
    ('absento (da-state-constraints-absents-pending-add st args))
    ('type (da-state-constraints-types-pending-add st args))
    (_ (error (format "unsupported constraint: ~a ~a" name args)))))
(def (da-constrain st)
  (values st vr-new) = (muk-sub-new-bindings st)
  dats = (da-constraints-new-bindings (muk-state-constraints st) vr-new)
  pending = (da-constraints-pending dats)
  dats =
  (if (null? pending) dats
    (lets
      (constraints ds as ts) = pending
      (values st var=>cxs) =
      (da-constrain-types st (da-constraints-var=>cxs dats) ts)
      (and var=>cxs
           (lets
             (values st var=>cxs) =
             (da-constrain-absents st var=>cxs as)
             (and var=>cxs
                  (forf dats = (da-constraints-var=>cxs-set
                                 (da-constraints-pending-clear dats) var=>cxs)
                        (list cadd constrain cx->var pending) <-
                        (list (list constraints-diseqs-add
                                    da-constrain-diseq
                                    diseq-cx->var
                                    ds))
                        #:break (not dats)
                        new = (let loop ((new '()) (pending pending))
                                (match pending
                                  ('() new)
                                  ((cons cx pending)
                                   (lets cxs = (constrain st cx)
                                         (and cxs (loop (append cxs new) pending))))))
                        (and new
                             (da-constraints-var=>cxs-set
                               dats
                               (forf var=>cxs = (da-constraints-var=>cxs dats)
                                     cx <- new
                                     (hash-update
                                       var=>cxs
                                       (cx->var cx)
                                       (lambda (cxs) (cadd cxs cx))
                                       constraints-empty-v))))))))))
  (or (and dats (list (muk-state-constraints-set st dats))) '()))

(define (da-constrain-diseq st or-diseqs)
  (def (muk-var< (muk-var n0) (muk-var n1)) (symbol<? n0 n1))
  (def (total< e0 e1)
    (or (not (muk-var? e1)) (and (muk-var? e0) (muk-var< e0 e1))))
  (def (cons< (cons k0 v0) (cons k1 v1)) (muk-var< k0 k1))
  (let ((st (forf st = st
                  (cons e0 e1) <- or-diseqs
                  #:break (not st)
                  (muk-unify st e0 e1))))
    (if st
      (lets (values st vr-new) = (muk-sub-new-bindings st)
            or-diseqs = (forl vr <- vr-new
                              (values _ val) = (muk-sub-get st vr)
                              (apply cons (sort (list vr val) total<)))
            or-diseqs = (sort or-diseqs cons<)
            (and (pair? or-diseqs) (list or-diseqs)))
      '())))
(define (=/=* or-diseqs) (muk-constraint '=/=* or-diseqs))
(define (=/= e0 e1) (=/=* `((,e0 . ,e1))))

(define (da-constrain-absents st var=>cxs as)
  (forf st = st var=>cxs = var=>cxs
        (list ground tm) <- as
        #:break (not var=>cxs)
        (letn loop (values st var=>cxs tm) = (values st var=>cxs tm)
          (values st tm) = (muk-walk st tm)
          (match tm
            ((? muk-var?)
             (values st (hash-update
                          var=>cxs tm
                          (lambda (cxs) (constraints-absents-add cxs ground))
                          constraints-empty-v)))
            ((cons h0 t0)
             (lets (values st var=>cxs) = (loop st var=>cxs h0)
                   (if var=>cxs
                     (loop st var=>cxs t0)
                     (values st var=>cxs))))
            (_ (values st (if (eqv? ground tm) #f var=>cxs)))))))
(define (absento ground tm)
  (if (pair? ground)
    (error (format "absento only supports absence of ground terms: ~a ~a"
                   ground tm))
    (muk-constraint 'absento (list ground tm))))

(define (da-constrain-types st var=>cxs ts)
  (forf st = st var=>cxs = var=>cxs
        (list tag tm) <- ts
        #:break (not var=>cxs)
        (values st tm) = (muk-walk st tm)
        (values st
          (if (muk-var? tm)
            (lets
              cxs = (hash-ref var=>cxs tm constraints-empty-v)
              tag0 = (constraints-type cxs)
              (if tag0
                (and (eq? tag0 tag) var=>cxs)
                (hash-set var=>cxs tm (constraints-type-set cxs tag))))
            (and (or (and (symbol? tm) (eq? tag 'sym))
                     (and (number? tm) (eq? tag 'num))) var=>cxs)))))
(define (typeo tag tm)
  (if (set-member? (set 'num 'sym) tag) (muk-constraint 'type (list tag tm))
    (error (format "invalid type tag: ~a ~a" tag tm))))
(define (symbolo tm) (typeo 'sym tm))
(define (numbero tm) (typeo 'num tm))
{'fraction_non_alphanumeric': 0.1412967103161929, 'fraction_numerical': 0.004854679016288725, 'mean_word_length': 1.8752984389348026, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 86, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 19, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2790462', 'n_tokens_mistral': 5324, 'n_tokens_neox': 5278, 'n_words': 1417}