;;; WebAssembly support library
;;; Copyright (C) 2023 David Thompson <dave@spritely.institute>
;;; Copyright (C) 2023, 2026 Igalia, S.L.
;;;
;;; 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:
;;;
;;; Validation of WebAssembly modules.
;;;
;;; Code:

(define-module (wasm validate)
  #:use-module (ice-9 binary-ports)
  #:use-module (ice-9 exceptions)
  #:use-module (ice-9 match)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-9 gnu)
  #:use-module (wasm parse)
  #:use-module (wasm stack)
  #:use-module (wasm canonical-types)
  #:use-module (wasm types)
  #:export (validate-wasm
            load-and-validate-wasm
            validated-wasm?
            validated-wasm-info
            validated-wasm-ref

            renumber-validated-wasm

            &wasm-error

            &wasm-validation-error
            wasm-validation-error?
            wasm-validation-error-module))


;;;
;;; Types
;;;

(define (s32? x)
  (and (exact-integer? x) (< (- (ash -1 31) 1) x (ash 1 31))))

(define (s64? x)
  (and (exact-integer? x) (< (- (ash -1 63) 1) x (ash 1 63))))

(define (f32? x)
  (and (number? x) (or (inexact? x) (exact-integer? x))))

(define (f64? x)
  (and (number? x) (or (inexact? x) (exact-integer? x))))

(define (resolve-type type)
  (match type
    (($ <sub-type> _ _ type) type)
    (_ type)))


;;;
;;; Validation
;;;

(define-record-type <validated-wasm>
  (make-validated-wasm info wasm)
  validated-wasm?
  (info validated-wasm-info)
  (wasm validated-wasm-ref))

(set-record-type-printer! <validated-wasm>
                          (lambda (x port)
                            (format port "#<validated-wasm ~a ~a>"
                                    (validated-wasm-ref x)
                                    (object-address x))))

(define-exception-type &wasm-error &error
  make-wasm-error
  wasm-error?)

(define-exception-type &wasm-validation-error &wasm-error
  make-wasm-validation-error
  wasm-validation-error?
  (wasm wasm-validation-error-wasm))

;; TODO: Trace instruction position within function to give context to
;; validation errors.
(define* (validate-wasm wasm #:key
                        (canonicalization (current-canonicalization)))
  (define wasm-info (make-wasm-info wasm canonicalization))
  (define elems (list->vector (wasm-elems wasm)))
  (define strings (list->vector (wasm-strings wasm)))
  (define (validation-error msg . irritants)
    (raise-exception
     (make-exception
      (make-wasm-validation-error wasm)
      (make-exception-with-message
       (format #f "WASM validation error: ~a" msg))
      (make-exception-with-irritants irritants))))
  (define (assert-s32 x)
    (unless (s32? x)
      (validation-error "i32 constant out of range" x)))
  (define (assert-s64 x)
    (unless (s64? x)
      (validation-error "i64 constant out of range" x)))
  (define (assert-f32 x)
    (unless (f32? x)
      (validation-error "f32 constant out of range" x)))
  (define (assert-f64 x)
    (unless (f64? x)
      (validation-error "f64 constant out of range" x)))
  (define-syntax-rule (define-lookup name proc)
    (define (name idx)
      (cdr (vector-ref (proc wasm-info) idx))))
  (define-lookup lookup-type wasm-info-types)
  (define-lookup lookup-global-type wasm-info-globals)
  (define-lookup lookup-tag wasm-info-tags)
  (define-syntax-rule (check-vector vec idx msg)
    (unless (< -1 idx (vector-length vec))
      (validation-error msg idx)))
  (define-syntax-rule (define-check name proc msg)
    (define (name idx)
      (check-vector (proc wasm-info) idx msg)))
  (define-check check-func wasm-info-funcs "invalid function")
  (define-check check-memory wasm-info-memories "invalid memory")
  (define-check check-table wasm-info-tables "invalid table")
  (define-check check-tag wasm-info-tags "invalid tag")
  (define (check-elem id) (check-vector elems id "invalid element"))
  (define (check-string id) (check-vector strings id "invalid string"))
  (define (assert-preceding-and-immutable-global idx current-idx)
    (if (< idx current-idx)
        (when (global-type-mutable? (lookup-global-type idx))
          (validation-error "mutable global reference in constant" idx))
        (validation-error "uninitialized global reference in constant" idx)))
  (define num-globals (vector-length (wasm-info-globals wasm-info)))
  (define (intern-heap-type ht)
    (match ht
      ((? symbol?) ht)
      ((? exact-integer?)
       (vector-ref (wasm-info-type-mapping wasm-info) ht))))
  (define (intern-ref-type rt)
    (match rt
      (($ <ref-type> nullable? ht)
       (make-ref-type nullable? (intern-heap-type ht)))))
  (define (intern-val-type type)
    (match type
      ((? symbol?) type) ; extern, i32, etc
      (_ (intern-ref-type type))))
  (define* (validate-const type instrs #:optional (global-count num-globals))
    (define (validate-instr ctx instr)
      (match (apply-stack-effect ctx (compute-stack-effect ctx instr))
        (($ <invalid-ctx> reason)
         (validation-error reason instr))
        (ctx
         (match instr
           (('i32.const x) (assert-s32 x))
           (('i64.const x) (assert-s64 x))
           (('f32.const x) (assert-f32 x))
           (('f64.const x) (assert-f64 x))
           (('global.get idx)
            (assert-preceding-and-immutable-global idx global-count))
           (('string.const idx) (check-string idx))
           (('ref.func f) (check-func f))
           ((or ('ref.null _)
                ('ref.i31)
                ('struct.new _)
                ('struct.new_default _)
                ('array.new _)
                ('array.new_default _)
                ('array.new_fixed _ _)
                ('any.convert_extern)
                ('extern.convert_any))
            #t)
           (_ (validation-error "invalid constant instruction" instr)))
         ctx)))
    ;; We need to make a phony func object that represents the
    ;; expected result type of the constant instructions.
    (let* ((sig (make-func-sig '() (list type)))
           (func (make-func #f (make-type-use #f sig) '() '())))
      (let loop ((ctx (initial-ctx wasm-info func))
                 (instrs instrs))
        (match instrs
          (() #t)
          ((instr . rest)
           (loop (validate-instr ctx instr) rest))))))
  (define (validate-global global idx)
    (match global
      (($ <global> id ($ <global-type> _ type) instrs)
       (validate-const type instrs idx))))
  (define (validate-data data)
    (match data
      (($ <data> _ _ _ offset _)
       (when offset
         (validate-const 'i32 offset)))))
  (define (validate-elem elem)
    (match elem
      (($ <elem> _ mode _ type offset inits)
       (when (eq? mode 'active)
         (validate-const 'i32 offset))
       (for-each (lambda (init)
                   (validate-const type init))
                 inits))))
  (define (validate-func func)
    (match func
      (($ <func> _ ($ <type-use> _ type) _ body)
       (define (lookup-block-type bt)
         (match bt
           (#f (make-func-sig '() '()))
           ((? exact-integer? idx) (lookup-type idx))
           (($ <type-use> #f
               ($ <func-sig>
                  (($ <param> _ (= intern-val-type ptype)) ...)
                  ((= intern-val-type rtype) ...)))
            (make-func-sig ptype rtype))
           (val-type
            (make-func-sig '() (list (intern-val-type val-type))))))
       (define (push-block* ctx type bt)
         (match bt
           (($ <func-sig> (($ <param> _ params) ...) results)
            (push-block ctx #f type params results))))
       (define (defaultable-type? type)
         (match type
           ((or 'i8 'i16 'i32 'i64 'f32 'f64) #t)
           (($ <ref-type> #t) #t)
           (_ #f)))
       (define (validate-instr ctx instr)
         (match (apply-stack-effect ctx (compute-stack-effect ctx instr))
           (($ <invalid-ctx> reason)
            (validation-error reason instr))
           (ctx
            (match instr
              (('if _ (= lookup-block-type bt) consequent alternate)
               (validate-branch (push-block* ctx 'if bt) consequent)
               (validate-branch (push-block* ctx 'if bt) alternate))
              (('block _ (= lookup-block-type bt) body)
               (validate-branch (push-block* ctx 'block bt) body))
              (('loop _ (= lookup-block-type bt) body)
               (validate-branch (push-block* ctx 'loop bt) body))
              (('try _ (= lookup-block-type bt) body catches catch-all)
               (match bt
                 (($ <func-sig> _ try-results)
                  (validate-branch (push-block* ctx 'try bt) body)
                  (for-each (match-lambda
                              ((tag-idx . body)
                               (check-tag tag-idx)
                               (match (lookup-tag tag-idx)
                                 (($ <tag-type>
                                     _ ($ <type-use>
                                          _ ($ <func-sig>
                                               (($ <param> _ params) ...)
                                               ())))
                                  (let ((ctx (push-block ctx #f 'catch
                                                         params
                                                         try-results)))
                                    (validate-branch ctx body)))
                                 (type
                                  (validation-error "invalid tag type" type)))))
                            catches)
                  (validate-branch (push-block ctx #f 'catch '() try-results)
                                   catch-all))))
              (('try_delegate _ (= lookup-block-type bt) body label)
               (validate-branch (push-block* ctx 'try bt) body))
              (('throw idx) (check-tag idx))
              (('rethrow idx)
               (match (lookup-block ctx idx)
                 (($ <block> _ 'catch) #t)
                 (($ <block> _ type)
                  (validation-error "rethrow of a non-catch block" idx type))))
              (('i32.const x) (assert-s32 x))
              (('i64.const x) (assert-s64 x))
              (('f32.const x) (assert-f32 x))
              (('f64.const x) (assert-f64 x))
              (('string.const idx) (check-string idx))
              (('global.set idx)
               (match (lookup-global-type idx)
                 (($ <global-type> mutable? _)
                  (unless mutable?
                    (validation-error "global is immutable" idx)))))
              (((or 'memory.size 'memory.grow) id)
               (check-memory id))
              (((or 'i32.load 'i32.load8_s 'i32.load8_u
                    'i32.load16_s 'i32.load16_u
                    'i64.load 'i64.load8_s 'i64.load8_u
                    'i64.load16_s 'i64.load16_u 'i64.load32_s 'i64.load32_u
                    'i32.store 'i32.store8 'i32.store16
                    'i64.store 'i64.store8 'i64.store16 'i64.store32)
                ($ <mem-arg> id _ _))
               (check-memory id))
              (((or 'table.set 'table.get 'table.size
                    'table.grow 'table.fill) table)
               (check-table table))
              (('table.copy dst src)
               (check-table dst)
               (check-table src))
              (('table.init table elem)
               (check-table table)
               (check-elem elem))
              (('elem.drop table) (check-elem table))
              (('ref.func f) (check-func f))
              (('struct.new_default t)
               ;; Validate that all fields have defaults.
               (match (resolve-type (lookup-type t))
                 ((and struct ($ <struct-type> (($ <field> _ _ types) ...)))
                  (unless (every defaultable-type? types)
                    (validation-error "struct type has non-defaultable fields"
                                      struct)))))
              (('struct.set t i)
               (let* ((type (resolve-type (lookup-type t)))
                      (field (list-ref (struct-type-fields type) i)))
                 (unless (field-mutable? field)
                   (validation-error "struct field is immutable" type field instr))))
              (('array.new_default t)
               (let ((array (resolve-type (lookup-type t))))
                 (unless (defaultable-type? (array-type-type array))
                   (validation-error "array type has non-defaultable element type"
                                     array))))
              ((or ('array.set t)
                   ('array.init_data t _)
                   ('array.init_elem t _))
               (let ((type (resolve-type (lookup-type t))))
                 (unless (array-type-mutable? type)
                   (validation-error "array is immutable" type instr))))
              ;; TODO: Validate array.get_s, ref.cast, etc.
              (_ #t))
            ctx)))
       (define (validate-branch ctx instrs)
         (let loop ((ctx ctx)
                    (instrs* instrs))
           (match instrs*
             (()
              (match (fallthrough ctx)
                (($ <invalid-ctx> reason)
                 (validation-error reason instrs))
                (_ #t)))
             ((instr . rest)
              (loop (validate-instr ctx instr) rest)))))
       (validate-branch (initial-ctx wasm-info func) body))))
  ;; Because global.get is a valid constant expression, we need to
  ;; track the index of the global that is currently being validated.
  ;; Any global.get instructions that reference an index greater than
  ;; or equal to the current global's index is invalid.
  (let ((imported-globals-count
         (fold (lambda (i sum)
                 (match i
                   (($ <import> _ _ 'global) (+ sum 1))
                   (_ sum)))
               0 (wasm-imports wasm))))
    (let loop ((globals (wasm-globals wasm))
               (i imported-globals-count))
      (match globals
        (() #t)
        ((global . rest)
         (validate-global global i)
         (loop rest (+ i 1))))))
  (for-each validate-func (wasm-funcs wasm))
  (for-each validate-data (wasm-datas wasm))
  (for-each validate-elem (wasm-elems wasm))
  (make-validated-wasm wasm-info wasm))

(define (renumber-validated-wasm wasm)
  (match wasm
    (($ <validated-wasm> info module)
     (make-validated-wasm (make-wasm-info
                           module
                           (renumber-canonicalization
                            (wasm-info-canonicalization info)))
                          module))))

(define* (load-and-validate-wasm obj #:key
                                 (canonicalization (current-canonicalization)))
  "Load and validate the WASM module within OBJ.  OBJ may be a <wasm>
record produced by 'resolve-wasm', a bytevector containing a WASM
binary, or an input port from which a WASM binary is read."
  (validate-wasm
   (cond
    ((wasm? obj) obj)
    ((bytevector? obj)
     (call-with-input-bytevector obj parse-wasm))
    ((port? obj)
     (parse-wasm obj))
    (else
     (error "not a WASM object" obj)))
   #:canonicalization canonicalization))
