;;; Vikalpa --- Prove S-expression
;;; Copyright © 2020 Masaya Tojo <masaya@tojo.tokyo>
;;;
;;; This file is part of Vikalpa.
;;;
;;; Vikalpa is free software; you can redistribute it and/or modify it
;;; under the terms of the GNU General Public License as published by
;;; the Free Software Foundation; either version 3 of the License, or
;;; (at your option) any later version.
;;;
;;; Vikalpa is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with Vikalpa.  If not, see <http://www.gnu.org/licenses/>.

(define-module (vikalpa)
  #:export (rewrite)
  #:use-module (vikalpa syntax)
  #:use-module (vikalpa primitive)
  #:use-module (ice-9 match)
  #:use-module ((srfi srfi-1) #:select (every any member))
  #:use-module (srfi srfi-8)
  #:use-module (srfi srfi-26))

(define (option p?)
  (lambda (x)
    (or (p? x) (not x))))

(define (list-of p?)
  (lambda (x)
    (and (list? x)
         (every p? x))))

;; (expression? x) -> boolean?
(define (expression? expr)
  (cond ((expr-quoted? expr)
         (or (integer? (expr-unquote expr))
             (boolean? (expr-unquote expr))
             (symbol? (expr-unquote expr))))
        ((pair? expr)
         (and (expression? (car expr))
              (list? (cdr expr))
              (every expression? (cdr expr))))
        ((variable? expr) #t)
        (else #f)))

(define (all-vars x)
  (cond ((expr-quoted? x) '())
        ((pair? x)
         (append (all-vars (car x))
                 (all-vars (cdr x))))
        ((variable? x) (list x))
        (else '())))

;; (expr-quoted? x) -> boolean?
(define (expr-quoted? expr)
  (and (pair? expr)
       (eq? 'quote (car expr))
       (pair? (cdr expr))
       ((const #t) (car (cdr expr)))
       (null? (cdr (cdr expr)))))

(define (expr-unquote expr)
  (car (cdr expr)))

;; (expr-quoted? x) -> boolean?
(define (expr-quote expr)
  (list 'quote expr))

;; (variable? x) -> boolean?
(define (variable? x)
  (symbol? x))

;; (variable=? variable? variable?) -> boolean?
(define (variable=? v1 v2)
  (eq? v1 v2))

;; (rule? x) -> boolean?
(define (rule? x)
  (and (list? x)
       (= 5 (length x))
       (eq? (car x) 'rule)
       (let ((vars (list-ref x 1)))
         (and (list? vars)
              (every variable? vars)))
       (let ((preconds (list-ref x 2)))
         (and (list? preconds)
              (every expression? preconds)))
       (expression? (list-ref x 3))
       (expression? (list-ref x 4))))

;; (rule (list-of variable?) expression? expression?) -> rule?
(define (rule vars preconds lhs rhs)
  (list 'rule vars preconds lhs rhs))

;; (rule-vars rule?) -> (list-of variable?)
(define (rule-vars r)
  (list-ref r 1))

;; (rule-preconds rule?) -> expression?
(define (rule-preconds r)
  (list-ref r 2))

;; (rule-lhs rule?) -> expression?
(define (rule-lhs r)
  (list-ref r 3))

;; (rule-rhs rule?) -> expression?
(define (rule-rhs r)
  (list-ref r 4))

;; (environment x) -> boolean?
;; environment is alist
(define (environment x)
  (cond ((null? x) #t)
        ((pair? x)
         (and (pair? (car x))
              (variable? (car (car x)))
              (expression? (cdr (car x)))
              (environment (cdr x))))
        (else #f)))

;; (apply-rule preconds rule? expression?) -> (or (cons rhs environment) #f)
(define (apply-rule preconds rl expr)
  (define (var? v)
    (and (member v (rule-vars rl))
         #t))
  (call/cc
   (lambda (k)
     (define (mat-map lhss exprs env)
       (cond ((and (pair? lhss)
                   (pair? exprs))
              (mat-map (cdr lhss)
                       (cdr exprs)
                       (mat (car lhss)
                            (car exprs)
                            env)))
             ((and (null? lhss)
                   (null? exprs))
              env)
             (else (k #f))))
     (define (mat lhs expr env)
       (cond ((expr-quoted? lhs)
              (if (expr-quoted? expr)
                  (and (equal? lhs expr)
                       env)
                  (k #f)))
             ((pair? lhs)
              (if (and (pair? expr)
                       (symbol? (car lhs))
                       (eqv? (car lhs) (car expr)))
                  (mat-map (cdr lhs) (cdr expr) env)
                  (k #f)))
             ((var? lhs)
              (cond
               ((assoc lhs env)
                => (match-lambda
                     ((env-var . env-expr)
                      (if (equal? env-expr expr)
                          env
                          (k #f)))))
               (else
                (cons (cons lhs expr)
                      env))))
             (else (k #f))))
     (define (mat-preconds rlps env)
       (if (null? rlps)
           env
           (mat-preconds (cdr rlps)
                         (any (lambda (p) (mat (car rlps) p env))
                              preconds))))
     (define (substitute env expr)
       (cond ((expr-quoted? expr) expr)
             ((pair? expr)
              (cons (substitute env (car expr))
                    (substitute env (cdr expr))))
             ((and (variable? expr) (assoc expr env)) => cdr)
             (else expr)))
     (substitute (mat (rule-lhs rl)
                      expr
                      (mat-preconds (rule-preconds rl) '()))
                 (rule-rhs rl)))))

;; (book? x) -> boolean?
;; book is alist
(define (book? x)
  (cond ((null? x) #t)
        ((pair? x)
         (and (pair? (car x))
              (or (proposition? (car x))
                  (primitive-function? (car x))
                  (macro? (car x))
                  (function? (car x)))
              (book? (cdr x))))
        (else #f)))

(define (primitive-function? x)
  (and (list? x)
       (= 4 (length x))
       (variable? (list-ref x 0))
       (eq? 'primitive-function (list-ref x 1))
       (let ((ps (list-ref x 2)))
         (and (list? ps)
              (every procedure? ps)))
       (procedure? (list-ref x 3))))

;; (primitive-function variable? (list-of procedure?) procedure?) -> primitive-function?
(define (primitive-function name ps proc)
  (list name 'primitive-function ps proc))

(define (primitive-function-name pf)
  (list-ref pf 0))

(define (primitive-function-preds pf)
  (list-ref pf 2))

(define (primitive-function-proc pf)
  (list-ref pf 3))

(define (match-primitive-function pf expression)
  (and (pair? expression)
       (eq? (primitive-function-name pf) (car expression))
       (list? (cdr expression))
       (every identity (map (lambda (pred expr)
                              (and (expr-quoted? expr)
                                   (pred (expr-unquote expr))))
                            (primitive-function-preds pf)
                            (cdr expression)))))

(define (apply-primitive-function pf expression)
  (apply (primitive-function-proc pf) (cdr expression)))

(define (function? x)
  (and (list? x)
       (= 4 (length x))
       (variable? (function-name x))
       (eq? 'function (list-ref x 1))
       (vars? (function-vars x))
       (expression? (function-expression x))))

;; (function variable? vars? expression?) -> function?
(define/guard (function (name variable?) (vars vars?) (expr expression?))
  (list name 'function vars expr))

(define (function-name f)
  (list-ref f 0))

(define (function-vars f)
  (list-ref f 2))

(define (function-expression f)
  (list-ref f 3))

(define (mrule? x)
  (and (list? x)
       (= 4 (length x))
       (eq? 'mrule (list-ref x 0))
       (vars? (mrule-vars x))
       ((const #t) (mrule-lhs x))
       ((const #t) (mrule-rhs x))))

(define (mrule vars lhs rhs)
  (list 'mrule vars lhs rhs))

(define (mrule-vars mrl)
  (list-ref mrl 1))

(define (mrule-lhs mrl)
  (list-ref mrl 2))

(define (mrule-rhs mrl)
  (list-ref mrl 3))

(define (macro? x)
  (and (list? x)
       (= 3 (length x))
       (eq? 'macro (list-ref x 1))
       (variable? (macro-name x))
       (let ((mrls (macro-mrules x)))
         (and (list? mrls)
              (every mrule? mrls)))))

(define (macro-name m)
  (list-ref m 0))

(define (macro-mrules m)
  (list-ref m 2))

(define (macro name rules)
  (list name 'macro rules))

(define (apply-mrule rl expr)
  (define (var? v)
    (and (member v (mrule-vars rl))
         #t))
  (call/cc
   (lambda (k)
     (define (mat-map lhs-tree expr-tree env)
       (cond ((and (pair? lhs-tree)
                   (pair? expr-tree))
              (mat-map (car lhs-tree)
                       (car expr-tree)
                       (mat-map (cdr lhs-tree)
                                (cdr expr-tree)
                                env)))
             (else (mat lhs-tree expr-tree env))))
     (define (mat lhs expr env)
       (cond ((expr-quoted? lhs)
              (if (expr-quoted? expr)
                  (and (equal? lhs expr)
                       env)
                  (k #f)))
             ((and (pair? lhs)
                   (pair? expr))
              (mat-map lhs expr env))
             ((var? lhs)
              (cond
               ((assoc lhs env)
                => (match-lambda
                     ((env-var . env-expr)
                      (if (equal? env-expr expr)
                          env
                          (k #f)))))
               (else
                (cons (cons lhs expr)
                      env))))
             ((eqv? lhs expr) env)
             (else (k #f))))
     (define (substitute env expr)
       (cond ((expr-quoted? expr) expr)
             ((pair? expr)
              (cons (substitute env (car expr))
                    (substitute env (cdr expr))))
             ((and (variable? expr) (assoc expr env)) => cdr)
             (else expr)))
     (substitute (mat (mrule-lhs rl)
                      expr
                      '())
                 (mrule-rhs rl)))))

(define (apply-macro m expr)
  (cond ((and (pair? expr)
              (eq? (macro-name m) (car expr)))
         (let loop ((rs (macro-mrules m)))
           (cond ((null? rs) expr)
                 ((apply-mrule (car rs) expr) => identity)
                 (else (loop (cdr rs))))))
        (else #f)))

(define/guard (expand (ms (list-of macro?)) (expr (const #t)))
  (let loop ((ms ms)
             (expr expr))
    (cond
     ((null? ms) expr)
     (else (expand (cdr ms)
                   (cond
                    ((apply-macro (car ms) expr) => identity)
                    (else expr)))))))

(define/guard (expand* (ms (list-of macro?)) (expr (const #t)))
  (let loop ((ms ms)
             (expr expr))
    (let ((new-expr (expand ms expr)))
      (if (equal? expr new-expr)
          (if (pair? new-expr)
              (cons (expand* ms (car new-expr))
                    (expand* ms (cdr new-expr)))
              new-expr)
          (expand* ms new-expr)))))

;; (axiom variable? vars? expression?) -> axiom?
;; axiom? is proposition
(define/guard (axiom (name variable?) (vars vars?) (expr expression?))
  (list name 'axiom vars expr))

(define (vars? x)
  (and (list? x)
       (every variable? x)))

;; (axiom? x) -> boolean?
(define (axiom? x)
  (and (list? x)
       (= 4 (length x))
       (variable?  (list-ref x 0))
       (eq? 'axiom (list-ref x 1))
       (vars? (list-ref x 2))
       (expression? (list-ref x 3))))

;; (theorem name vars? expression?) -> theorem?
(define/guard (theorem (name variable?) (vars vars?) (expr expression?))
  (list name 'theorem vars expr))

;; (theorem? x) -> boolean?
;; theorem? is proposition
(define (theorem? x)
  (and (list? x)
       (= 4 (length x))
       (variable? (list-ref x 0))
       (eq? 'theorem (list-ref x 1))
       (vars? (list-ref x 2))
       (expression? (list-ref x 3))))

(define (theorem-name x)
  (list-ref x 0))

;; (proposition? x) -> boolean?
(define (proposition? x)
  (or (axiom? x)
      (theorem? x)))

(define (proposition-name x)
  (list-ref x 0))

(define (proposition-vars x)
  (list-ref x 2))

(define (proposition-expression x)
  (list-ref x 3))

(define (proposition-rules x)
  (expression->rules (proposition-vars x)
                     '()
                     (proposition-expression x)))

;; (rewriter path? command?) -> rewriter?
(define (rewriter p c)
  (list p c))

(define (rewriter-path r)
  (list-ref r 0))

(define (rewriter-command r)
  (list-ref r 1))

(define (rewriter? x)
  (and (pair? x)
       (path? (car x))
       (pair? (cdr x))
       (command? (car (cdr x)))
       (null? (cdr (cdr x)))))

;; (sequence? x) -> boolean?
(define (sequence? x)
  (and (list? x)
       (every rewriter? x)))

;; (path? x) -> boolean?
;; path is list
(define (path? x)
  (and (list? x)
       (every natural? x)))

;; (command? x) -> booelan?
(define (command? x)
  (or (symbol? x)
      ((list-of symbol?) x)))

;; (extract path? expression? preconds? procedure?) -> (values expression? procedure?)
(define (extract path expr fail)
  (if (null? path)
      (values expr '() identity)
      (let ((i (car path)))
        (cond
         ((expr-if? expr)
          (let ((precond (expr-if-cond expr)))
            (receive (extracted-expr extracted-preconds builder)
                (extract (cdr path) (list-ref expr i) fail)
              (values extracted-expr
                      (case i
                        ((1) '())
                        ((2) (cons precond extracted-preconds))
                        ((3) (cons (expr-not precond) extracted-preconds))
                        (else (fail "(if) invaild path" path)))
                      (lambda (x)
                        (append (list-head expr i)
                                (list (builder x))
                                (list-tail expr (+ i 1))))))))
         ((< i (length expr))
          (receive (extracted-expr preconds builder)
              (extract (cdr path) (list-ref expr i) fail)
            (values extracted-expr
                    preconds
                    (lambda (x)
                      (append (list-head expr i)
                              (list (builder x))
                              (list-tail expr (+ i 1)))))))
         (else
          (fail "invalid path" path))))))

(define (function->rules x)
  (list (rule (function-vars x)
              '()
              (cons (function-name x)
                    (function-vars x))
              (function-expression x))
        (rule (function-vars x)
              '()
              (function-expression x)
              (cons (function-name x)
                    (function-vars x)))))

;; (rewrite book? rewriter? expression? procedure?) -> expr
(define (rewrite1 b expr fail r)
  (let ((cmd (rewriter-command r)))
    (receive (extracted-expr preconds builder)
        (extract (rewriter-path r) expr fail)
      (builder
       (cond
        ((equal? cmd '(expand))
         (expand (filter macro? b) extracted-expr))
        ((equal? cmd '(expand*))
         (expand* (filter macro? b) extracted-expr))
        ((and (symbol? cmd)
              (assoc cmd b))
         => (lambda (x)
              (cond
               ((proposition? x)
                (cond
                 ((any (cut apply-rule preconds <> extracted-expr)
                       (proposition-rules x)) => identity)
                 (else
                  (fail (format #f "no match-rule (~a)" cmd)))))
               ((function? x)
                (cond
                 ((any (cut apply-rule '() <> extracted-expr)
                       (function->rules x)) => identity)
                 (else
                  (fail (format #f "no match-rule (~a)" cmd)))))
               ((primitive-function? x)
                (cond
                 ((match-primitive-function x extracted-expr)
                  (apply-primitive-function x extracted-expr))
                 (else
                  (fail "can't apply-primitive-function (~a)" cmd))))
               ((macro? x)
                (cond
                 ((apply-macro x extracted-expr) => identity)
                 (else
                  (fail "macro fail"))))
               (else
                (fail "rewrite error")))))
        (else (fail "no cmd (~a)" cmd)))))))

(define/guard (rewrite (b book?) (expr expression?) (seq sequence?))
  (let loop ((expr expr)
             (seq seq))
    (if (null? seq)
        expr
        (loop (rewrite1 b
                        expr
                        (lambda (x . args)
                          (apply error
                                 (format #f "rewrite: ~a" x)
                                 args))
                        (car seq))
              (cdr seq)))))

(define (expr-not x)
  (list 'not x))

(define (expr-equal? x)
  (and (list? x)
       (= 3 (length x))
       (eq? 'equal? (list-ref x 0))
       (expression? (list-ref x 1))
       (expression? (list-ref x 2))))

(define (expr-equal-lhs x)
  (list-ref x 1))

(define (expr-equal-rhs x)
  (list-ref x 2))

(define (expr-if? x)
  (and (list? x)
       (= 4 (length x))
       (eq? 'if (list-ref x 0))
       (expression? (list-ref x 1))
       (expression? (list-ref x 2))
       (expression? (list-ref x 3))))

(define (expr-if-cond x)
  (list-ref x 1))

(define (expr-if-then x)
  (list-ref x 2))

(define (expr-if-else x)
  (list-ref x 3))

(define (expression->rules vars preconds expr)
  (if (expr-if? expr)
      (append (expression->rules vars
                                 (cons (expr-if-cond expr) preconds)
                                 (expr-if-then expr))
              (expression->rules vars
                                 (cons (expr-not (expr-if-cond expr)) preconds)
                                 (expr-if-else expr)))
      (if (expr-equal? expr)
          (list (rule vars
                      preconds
                      (expr-equal-lhs expr)
                      (expr-equal-rhs expr))
                (rule vars
                      preconds
                      (expr-equal-rhs expr)
                      (expr-equal-lhs expr)))
          '())))

(define (proposition->rules def)
  (expression->rules (proposition-vars def)
                     '()
                     (proposition-expression def)))

(define current-book (make-parameter '()))

(define (add-book x)
  (let ((b (current-book)))
    (current-book
     (cond ((assoc (proposition-name x) b)
            => (lambda (prop)
                 (if (equal? x prop)
                     b
                     (error "add-book: error"))))
           (else (cons x b))))))

(define-syntax define-axiom
  (syntax-rules ()
    ((_ name (var ...) expr)
     (add-book (axiom 'name '(var ...) 'expr)))))

(define-syntax define-theorem
  (syntax-rules ()
    ((_ name (var ...) expr)
     (add-book (theorem 'name '(var ...) 'expr)))))

(define-syntax define-function
  (syntax-rules ()
    ((_ name (var ...) expr)
     (let ((f (function 'name '(var ...) 'expr)))
       (add-book f)
       f))))

(define-syntax define-macro
  (syntax-rules ()
    ((_ name () ((lhs1 . lhs2) rhs) ...)
     (add-book (macro 'name
                 (list (mrule (all-vars 'lhs2) '(lhs1 . lhs2) 'rhs)
                       ...))))))

(define-syntax define-book
  (syntax-rules ()
    ((_ name (books ...) expr ...)
     (define* (name #:optional (parent '()))
       (parameterize ([current-book ((compose books ... identity) parent)])
         expr
         ...
         (unless (book? (current-book))
           (error "define-book: error"))
         (current-book))))))

(define-book prelude ()
  (define-axiom equal-same (x)
    (equal? (equal? x x) '#t))
  (define-axiom equal-swap (x y)
    (equal? (equal? x y) (equal? y x)))
  (define-axiom equal-if (x y)
    (if (equal? x y) (equal? x y) '#t))
  (define-axiom pair/cons (x y)
    (equal? (pair? (cons x y)) '#t))
  (define-axiom car/cons (x y)
    (equal? (car (cons x y)) x))
  (define-axiom cdr/cons (x y)
    (equal? (cdr (cons x y)) y))
  (define-axiom cons/car+cdr (x y)
    (equal? (if (pair? x)
                (equal? (cons (car x) (cdr y)) x)
                '#t)
            y))
  (define-axiom cons/car+cdr (x y)
    (equal? (if (pair? x)
                (equal? (cons (car x) (cdr y)) x)
                '#t)
            y))
  (define-axiom if-same (x y)
    (equal? (if x y y) y))
  (define-axiom size/car (x)
    (if (pair? x)
        (equal? (< (size x) (size (car x)))
                '#t)
        '#t))
  (define-axiom size/cdr (x)
    (if (pair? x)
        (equal? (< (size x) (size (cdr x)))
                '#t)
        '#t))
  (define-macro list ()
    ((list) '())
    ((list x . y) (cons x (list . y))))
  (define-macro + ()
    ((+) '0)
    ((+ x) x)
    ((+ x . xs) (b+ x (+ . xs))))
  (define-macro - ()
    ((- x) (u- x))
    ((- x . xs) (+ x . (u- (+ . xs)))))
  (define-macro and ()
    ((and) '#t)
    ((and x) x)
    ((and x . xs) (if x (and . xs) #f)))
  (define-macro or ()
    ((or) '#f)
    ((or x) x)
    ((or x . xs) (if x x (or . xs))))
  (define-macro atom? ()
    ((atom? x) (not (pair? x))))
  (define-axiom if-not (x y z)
    (equal? (if (not x)
                y
                z)
            (if x
                z
                y))))

(define-book app (prelude)
  (define-function app (x y)
    (if (atom? x)
        y
        (cons (car x) (app (cdr x) y))))

  (define-theorem associate-app (x y z)
    (equal? (app (app x y) z)
            (app x (app y z)))))

(define (size x)
  (if (pair? x)
      (+ 1
         (size (car x))
         (size (cdr x)))
      0))

(define (atom? x)
  (not (pair? x)))

(define (app x y)
    (if (atom? x)
        y
        (cons (car x) (app (cdr x) y))))