From 31d601237b14745e3353d91bacdca61bc79d367d Mon Sep 17 00:00:00 2001
From: Stas Boukarev <stassats@gmail.com>
Date: Wed, 1 May 2024 21:33:29 +0300
Subject: [PATCH] Don't disable argument checking with multiple callers.

---
 src/compiler/callable-args.lisp |  7 +++++++
 src/compiler/ir1util.lisp       | 22 ++++++++++++++++++++++
 tests/compiler-2.pure.lisp      | 11 +++++++++++
 3 files changed, 40 insertions(+)

diff --git a/src/compiler/callable-args.lisp b/src/compiler/callable-args.lisp
index e8efa6fc3..7b4f55d0e 100644
--- a/src/compiler/callable-args.lisp
+++ b/src/compiler/callable-args.lisp
@@ -498,6 +498,13 @@
 
 (defun disable-arg-count-checking (leaf type arg-count)
   (when (lambda-p leaf)
+    (let ((once nil))
+      ;; TODO: what if all destinations can disable arg count checking.
+      (map-leaf-refs (lambda (dest)
+                       (declare (ignore dest))
+                       (when (shiftf once t)
+                         (return-from disable-arg-count-checking)))
+                     leaf))
     (multiple-value-bind (min max) (fun-type-arg-limits type)
       (when (and min
                  (if max
diff --git a/src/compiler/ir1util.lisp b/src/compiler/ir1util.lisp
index 15a3e6200..37e636f64 100644
--- a/src/compiler/ir1util.lisp
+++ b/src/compiler/ir1util.lisp
@@ -3055,6 +3055,28 @@ is :ANY, the function name is not checked."
                   when (eq v lambda-var)
                   do (funcall function combination arg))))))))
 
+(defun map-leaf-refs (function leaf)
+  (let* ((seen-calls))
+    (labels ((recur (leaf)
+               (dolist (ref (leaf-refs leaf))
+                 (let* ((lvar (node-lvar ref))
+                        (dest (and lvar
+                                   (lvar-dest lvar))))
+                   (if (and (combination-p dest)
+                            (eq (combination-kind dest) :local))
+                       (let ((lambda (combination-lambda dest)))
+                         (when (cond ((functional-kind-eq lambda let))
+                                     ((memq dest seen-calls)
+                                      nil)
+                                     (t
+                                      (push dest seen-calls)))
+                           (loop for v in (lambda-vars lambda)
+                                 for arg in (combination-args dest)
+                                 when (eq arg lvar)
+                                 do (recur v))))
+                       (funcall function dest))))))
+      (recur leaf))))
+
 (defun propagate-lvar-annotations-to-refs (lvar var)
   (when (lvar-annotations lvar)
     (dolist (ref (leaf-refs var))
diff --git a/tests/compiler-2.pure.lisp b/tests/compiler-2.pure.lisp
index e806c6824..6e9774289 100644
--- a/tests/compiler-2.pure.lisp
+++ b/tests/compiler-2.pure.lisp
@@ -4404,3 +4404,14 @@
       (declare ((function (fixnum &rest t)) j))
       (apply j l r))
    ((#'+ 1 '(2)) 3)))
+
+(with-test (:name :disabling-arg-count-checking)
+  (checked-compile-and-assert
+      (:optimize :safe)
+      `(lambda (x d)
+         (let ((f (lambda (x y)
+                    (< x y))))
+           (funcall d f)
+           (sort x f)))
+    ((nil #'funcall) (condition 'program-error))
+    ((nil #'list) nil)))
-- 
2.11.4.GIT