;;; WebAssembly VM
;;; Copyright (C) 2026 Igalia, S.L.
;;; Copyright (C) 2023 David Thompson <dave@spritely.institute>
;;;
;;; Licensed under the Apache License, Version 2.0 (the "License");
;;; you may not use this file except in compliance with the License.
;;; You may obtain a copy of the License at
;;;
;;;    http://www.apache.org/licenses/LICENSE-2.0
;;;
;;; Unless required by applicable law or agreed to in writing, software
;;; distributed under the License is distributed on an "AS IS" BASIS,
;;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;;; See the License for the specific language governing permissions and
;;; limitations under the License.

;;; Commentary:
;;;
;;; Reference type canonicalization.
;;;
;;; Code:

(define-module (wasm canonical-types)
  #:use-module (ice-9 match)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (wasm types)
  #:export (make-canonicalization
            current-canonicalization
            canonical-type-ref
            canonicalize-types!
            canonical-types
            canonical-type-ref
            canonical-mutable-i8-array-heap-type
            canonical-mutable-i16-array-heap-type

            renumber-canonicalization
            canonical-type-bounds))

(define-record-type <growable-array>
  (%make-growable-array length blocks)
  growable-array?
  (length growable-array-length set-growable-array-length!)
  (blocks growable-array-blocks))

(define* (make-growable-array #:optional (size 8))
  (%make-growable-array 0 (list (make-vector size #f))))

(define (growable-array-ref array idx)
  (unless (< idx (growable-array-length array))
    (error "out of bounds" idx))
  (let lp ((blocks (growable-array-blocks array)) (idx idx))
    (match blocks
      ((block . blocks)
       (if (< idx (vector-length block))
           (vector-ref block idx)
           (lp blocks (- idx (vector-length block))))))))

(define (growable-array-push! array elt)
  (let ((length (growable-array-length array)))
    (let lp ((blocks (growable-array-blocks array)) (idx length))
      (match blocks
        ((block)
         (if (< idx (vector-length block))
             (vector-set! block idx elt)
             (lp (let ((tail (list (make-vector length #f))))
                   (set-cdr! blocks tail)
                   tail)
                 (- idx (vector-length block)))))
        ((block . blocks)
         (lp blocks (- idx (vector-length block))))))
    (set-growable-array-length! array (1+ length))
    length))

(define (growable-array->vector array)
  (let ((ret (make-vector (growable-array-length array) #f)))
    (let lp ((blocks (growable-array-blocks array)) (idx 0))
      (match blocks
        ((block)
         (vector-copy! ret idx block 0 (- (vector-length ret) idx))
         ret)
        ((block . blocks)
         (vector-copy! ret idx block)
         (lp blocks (+ idx (vector-length block))))))))

(define well-known-i8-array 0)
(define well-known-i16-array 1)
(define well-known-type-count 2)
(define (make-well-known-types) (make-vector well-known-type-count #f))

(define-record-type <canonicalization>
  (%make-canonicalization groups types well-known-types
                          bounds
                          any-bounds eq-bounds
                          struct-bounds array-bounds func-bounds)
  canonicalization?
  (groups canonicalization-groups)
  (types canonicalization-types)
  (well-known-types canonicalization-well-known-types)
  (bounds canonicalization-type-bounds)
  (any-bounds canonicalization-any-bounds
              set-canonicalization-any-bounds!)
  (eq-bounds canonicalization-eq-bounds
             set-canonicalization-eq-bounds!)
  (struct-bounds canonicalization-struct-bounds
                 set-canonicalization-struct-bounds!)
  (array-bounds canonicalization-array-bounds
                set-canonicalization-array-bounds!)
  (func-bounds canonicalization-func-bounds
               set-canonicalization-func-bounds!))

(define (make-canonicalization)
  (let ((groups (make-hash-table))
        (types (make-growable-array))
        (well-known (make-well-known-types))
        (type-bounds #())
        (empty-bounds '(0 . 0)))
    (%make-canonicalization groups types well-known type-bounds
                            ;; For any, eq, struct, array, and func type
                            ;; bounds, initially this is the empty set.
                            empty-bounds empty-bounds empty-bounds
                            empty-bounds empty-bounds)))

(define current-canonicalization
  (make-parameter (make-canonicalization)))

(define (roll type module-canonical-types)
  (define ntypes (growable-array-length module-canonical-types))
  (define (roll-heap-type type)
    (match type
      ;; Maybe it's something like 'extern.  Should error though if it's a
      ;; $name.
      ((? symbol?) type)
      ((? exact-integer?)
       ;; If the heap type is defined before this group, rewrite it to a
       ;; canonical type reference.  Otherwise make it negative, to
       ;; indicate that it's a recursive reference.
       (if (< type ntypes)
           (growable-array-ref module-canonical-types type)
           (- -1 (- type ntypes))))))

  (define (roll-value-type type)
    (match type
      ((? symbol?) type)
      (($ <ref-type> nullable? heap-type)
       (make-ref-type nullable? (roll-heap-type heap-type)))))

  (define (roll-base-type type)
    (match type
      (($ <func-sig> params results)
       (make-func-sig (map (match-lambda
                             (($ <param> id type)
                              (make-param #f (roll-value-type type))))
                           params)
                      (map roll-value-type results)))
      (($ <struct-type> fields)
       (make-struct-type
        (map (match-lambda
               (($ <field> id mutable? type)
                (make-field #f mutable? (roll-value-type type))))
             fields)))
      (($ <array-type> mutable? type)
       (make-array-type mutable? (roll-value-type type)))))

  (define (roll-sub-type type)
    (match type
      (($ <sub-type> final? supers type)
       (make-sub-type final?
                      (map roll-heap-type supers)
                      (roll-base-type type)))
      (_
       (roll-base-type type))))

  (roll-sub-type type))

;; Unrolling a type replaces relative recursive type indices with
;; canonical type references.
(define (unroll type canonical-group-start)
  (define (unroll-heap-type type)
    (match type
      ((? symbol?) type)
      ((? exact-integer? idx)
       (if (negative? idx)
           (+ canonical-group-start (- -1 idx))
           idx))))

  (define (unroll-value-type type)
    (match type
      ((? symbol?) type)
      (($ <ref-type> nullable? heap-type)
       (make-ref-type nullable? (unroll-heap-type heap-type)))))

  (define (unroll-base-type type)
    (match type
      (($ <func-sig> params results)
       (make-func-sig (map (match-lambda
                             (($ <param> _ type)
                              (make-param #f (unroll-value-type type))))
                           params)
                      (map unroll-value-type results)))
      (($ <struct-type> fields)
       (make-struct-type
        (map (match-lambda
               (($ <field> _ mutable? type)
                (make-field #f mutable? (unroll-value-type type))))
             fields)))
      (($ <array-type> mutable? type)
       (make-array-type mutable? (unroll-value-type type)))))

  (define (unroll-sub-type type)
    (match type
      (($ <sub-type> final? supers type)
       (make-sub-type final?
                      (map unroll-heap-type supers)
                      (unroll-base-type type)))
      (_ (unroll-base-type type))))

  (unroll-sub-type type))

(define (canonicalize-types! canonicalization types)
  "Canonicalize @var{types}, which should be a list of @code{<type>} or
@code{<rec-group>} terms, as from the @code{types} field of a
@code{<wasm>} object.

The types are added to @var{canonicalization}, which is a set of
canonical types shared across a number of wasm modules; if a term
already has a structurally compatible definition in the array,
canonicalizing it will replace it with a reference to the existing
canonicalized definition, and otherwise will add it to the
canonicalization.  Note that field, parameter, and heap type names are
ignored when canonicalizing types; a canonical type includes no names.

Return a vector mapping indexes into @var{types} (\"local type
indices\") to indexes into the canonicalization (\"canonical type
indices\").  Use @code{canonical-type-ref} to get the type object
corresponding to a canonical type index."
  (define canonical-groups (canonicalization-groups canonicalization))
  (define canonical-types (canonicalization-types canonicalization))
  ;; Local type index -> canonical type index.
  (define module-canonical-types (make-growable-array))

  (define (invalidate-computed-type-bounds! type)
    (define (invalidate-base-type! type)
      (match type
        (($ <struct-type>)
         (set-canonicalization-struct-bounds! canonicalization #f)
         (set-canonicalization-eq-bounds! canonicalization #f)
         (set-canonicalization-any-bounds! canonicalization #f))
        (($ <array-type>)
         (set-canonicalization-array-bounds! canonicalization #f)
         (set-canonicalization-eq-bounds! canonicalization #f)
         (set-canonicalization-any-bounds! canonicalization #f))
        (($ <func-sig>)
         (set-canonicalization-func-bounds! canonicalization #f))))
    (match type
      (($ <sub-type> final? () base-type)
       (invalidate-base-type! base-type))
      (($ <sub-type> final? (super) base-type)
       (let ((bounds (canonicalization-type-bounds canonicalization)))
         (when (and (< super (vector-length bounds))
                    (vector-ref bounds super))
           (vector-set! bounds super #f)
           (invalidate-computed-type-bounds!
            (growable-array-ref canonical-types super)))))
      (_ (invalidate-base-type! type))))

  (define (intern-type! type)
    (invalidate-computed-type-bounds! type)
    (growable-array-push! canonical-types type))

  (define (visit-group types)
    ;; Rolling up a type replaces indices outside of the type group
    ;; with a canonical type descriptor and indices inside of the
    ;; type group with relative indices.  This creates a new type
    ;; that can be equal? tested against a type from another module.

    ;; If a type group with identical structure has already been
    ;; canonicalized, return the cached type descriptors.  Otherwise,
    ;; generate new ones, cache them, and return them.
    (let ((rolled (map (lambda (type)
                         (roll type module-canonical-types))
                       types)))
      (cond
       ((hash-ref canonical-groups rolled)
        ;; Cache hit: Just copy 'em over.
        => (lambda (cached-group)
             (for-each
              (lambda (canonical-idx)
                (growable-array-push! module-canonical-types canonical-idx))
              cached-group)))
       ;; Cache miss: Generate and cache new descriptors.
       (else
        (let ((group-start (growable-array-length canonical-types)))
          (hash-set!
           canonical-groups
           rolled
           (map-in-order
            (lambda (rolled-type)
              (define canonical-idx
                (intern-type! (unroll rolled-type group-start)))
              (growable-array-push! module-canonical-types canonical-idx)
              canonical-idx)
            rolled)))))))

  ;; Visit all the type groups and canonicalize them.  A type that
  ;; is not in a recursive type group is treated as being in a group
  ;; of one.
  (for-each (lambda (group)
              (match group
                (($ <rec-group> (($ <type> id type) ...))
                 (visit-group type))
                (($ <type> id type)
                 (visit-group (list type)))))
            types)

  (growable-array->vector module-canonical-types))

(define (canonical-types canonicalization)
  (growable-array->vector (canonicalization-types canonicalization)))

(define (canonical-type-ref canonicalization idx)
  (growable-array-ref (canonicalization-types canonicalization) idx))

(define (intern-well-known-type! canonicalization well-known-types idx type
                                 ensure?)
  (define (intern! ht)
    (vector-set! well-known-types idx ht)
    ht)
  (if ensure?
      (match (canonicalize-types! canonicalization (list (make-type #f type)))
        (#(ht) (intern! ht)))
      (match (hash-ref (canonicalization-groups canonicalization) (list type))
        ((ht) (intern! ht))
        (#f #f))))

(define-syntax-rule (define-well-known-type getter idx type)
  (define* (getter canonicalization #:key ensure?)
    (let ((well-known (canonicalization-well-known-types canonicalization)))
      (or (vector-ref well-known idx)
          (intern-well-known-type! canonicalization well-known idx type
                                   ensure?)))))

(define-well-known-type canonical-mutable-i8-array-heap-type
  well-known-i8-array
  (make-array-type #t 'i8))
(define-well-known-type canonical-mutable-i16-array-heap-type
  well-known-i16-array
  (make-array-type #t 'i16))

(define (renumber-canonicalization canonicalization)
  "Sort the heap type values assigned to types in @var{canonicalization},
and return a new canonicalization.

A canonicalization defines a mapping between heap type
indices (integers) and heap type definitions (array, struct, and func
types).  Renumbering a canonicalization sorts the heap type indices such
that subtypes are contiguous.  This allows subtyping checks to just
compare a value's type tag to a lower and upper bound.

If a fresh type definition is later added to the renumbered
canonicalization, the computed bounds on that type's supertypes will be
invalidated."
  (define struct-roots '())
  (define array-roots '())
  (define func-roots '())

  (define types (canonical-types canonicalization))
  (define edges (make-vector (vector-length types) '()))
  (let lp ((i 0))
    (define (visit-base-type type)
      (match type
        (($ <struct-type>)
         (set! struct-roots (cons i struct-roots)))
        (($ <array-type>)
         (set! array-roots (cons i array-roots)))
        (($ <func-sig>)
         (set! func-roots (cons i func-roots)))))
    (when (< i (vector-length types))
      (match (vector-ref types i)
        (($ <sub-type> final? () base-type)
         (visit-base-type base-type))
        (($ <sub-type> final? (super) base-type)
         (vector-set! edges super (cons i (vector-ref edges super))))
        (base-type
         (visit-base-type base-type)))
      (lp (1+ i))))

  ;; Old (canonicalized) type index -> new (sorted) type index.
  (define renumbering (make-vector (vector-length types) #f))
  ;; New type index -> index of next type that is not a subtype.
  (define subtype-bounds (make-vector (vector-length types) #f))
  (define (compute-renumbering! start subtypes)
    (fold-right
     (lambda (subtype subtype-begin)
       (define subtype-end
         (compute-renumbering! (1+ subtype-begin) (vector-ref edges subtype)))
       (vector-set! renumbering subtype subtype-begin)
       (vector-set! subtype-bounds subtype-begin
                    (cons subtype-begin subtype-end))
       subtype-end)
     start
     subtypes))

  (define struct-start 0)
  (define struct-end (compute-renumbering! struct-start struct-roots))
  (define array-start struct-end)
  (define array-end (compute-renumbering! array-start array-roots))
  (define func-start array-end)
  (define func-end (compute-renumbering! func-start func-roots))

  (define (renumber-heap-type type)
    (match type
      ((? symbol?) type)
      ((? exact-integer?)
       (if (negative? type)
           type
           (vector-ref renumbering type)))))
  (define (renumber-value-type type)
    (match type
      ((? symbol?) type)
      (($ <ref-type> nullable? heap-type)
       (make-ref-type nullable? (renumber-heap-type heap-type)))))
  (define (renumber-base-type type)
    (match type
      (($ <func-sig> params results)
       (make-func-sig (map (match-lambda
                             (($ <param> id type)
                              (make-param id (renumber-value-type type))))
                           params)
                      (map renumber-value-type results)))
      (($ <struct-type> fields)
       (make-struct-type
        (map (match-lambda
               (($ <field> id mutable? type)
                (make-field id mutable? (renumber-value-type type))))
             fields)))
      (($ <array-type> mutable? type)
       (make-array-type mutable? (renumber-value-type type)))))
  (define (renumber-sub-type type)
    (match type
      (($ <sub-type> final? supers type)
       (make-sub-type final?
                      (map renumber-heap-type supers)
                      (renumber-base-type type)))
      (_
       (renumber-base-type type))))

  (define renumbered-types (make-vector (vector-length types) #f))
  (let lp ((i 0))
    (when (< i (vector-length types))
      (vector-set! renumbered-types (vector-ref renumbering i)
                   (renumber-sub-type (vector-ref types i)))
      (lp (1+ i))))

  (let ((groups (make-hash-table))
        (types (make-growable-array (vector-length renumbered-types))))
    (hash-for-each (match-lambda*
                     (((type ...) (idx ...))
                      (hash-set! groups (map renumber-sub-type type)
                                 (map renumber-heap-type idx))))
                   (canonicalization-groups canonicalization))

    (let lp ((i 0))
      (when (< i (vector-length renumbered-types))
        (growable-array-push! types (vector-ref renumbered-types i))
        (lp (1+ i))))

    (%make-canonicalization groups types (make-well-known-types)
                            subtype-bounds
                            (cons struct-start array-end) ; any
                            (cons struct-start array-end) ; eq
                            (cons struct-start struct-end); struct
                            (cons array-start array-end)  ; array
                            (cons func-start func-end)))) ; func

(define (canonical-type-bounds canonicalization heap-type)
  (match heap-type
    ((? exact-integer?)
     (let ((bounds (canonicalization-type-bounds canonicalization)))
       (and (< heap-type (vector-length bounds))
            (vector-ref bounds heap-type))))
    ('any (canonicalization-any-bounds canonicalization))
    ('eq (canonicalization-eq-bounds canonicalization))
    ('struct (canonicalization-struct-bounds canonicalization))
    ('array (canonicalization-array-bounds canonicalization))
    ('func (canonicalization-func-bounds canonicalization))))
