From a6f3e15cec7d4ef42d1ba82e42fe6d6656d09422 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Mon, 28 Jun 2021 03:34:05 -0700
Subject: [PATCH] [sledge] Add List.fold_diagonal

Reviewed By: ngorogiannis

Differential Revision: D28907802

fbshipit-source-id: 57b7b81c0
---
 sledge/nonstdlib/list.ml  | 9 +++++++++
 sledge/nonstdlib/list.mli | 1 +
 2 files changed, 10 insertions(+)

diff --git a/sledge/nonstdlib/list.ml b/sledge/nonstdlib/list.ml
index 36b1812cf..e236b797e 100644
--- a/sledge/nonstdlib/list.ml
+++ b/sledge/nonstdlib/list.ml
@@ -61,6 +61,15 @@ let fold_right xs init ~f = fold_right ~f ~init xs
 let reduce xs ~f =
   match xs with [] -> None | x :: xs -> Some (fold ~f xs x)
 
+let fold_diagonal xs init ~f =
+  let rec fold_diagonal_ z = function
+    | [] -> z
+    | x :: ys ->
+        let z = fold ~f:(fun y z -> f x y z) ys z in
+        fold_diagonal_ z ys
+  in
+  fold_diagonal_ init xs
+
 let fold_map xs init ~f =
   Pair.swap (fold_map ~f:(fun s x -> Pair.swap (f x s)) ~init xs)
 
diff --git a/sledge/nonstdlib/list.mli b/sledge/nonstdlib/list.mli
index 8d2782aad..24d12046a 100644
--- a/sledge/nonstdlib/list.mli
+++ b/sledge/nonstdlib/list.mli
@@ -95,6 +95,7 @@ val fold : 'a list -> 's -> f:('a -> 's -> 's) -> 's
 val fold_left : 'a list -> 's -> f:('s -> 'a -> 's) -> 's
 val fold_right : 'a list -> 's -> f:('a -> 's -> 's) -> 's
 val reduce : 'a t -> f:('a -> 'a -> 'a) -> 'a option
+val fold_diagonal : 'a list -> 's -> f:('a -> 'a -> 's -> 's) -> 's
 val fold_map : 'a t -> 's -> f:('a -> 's -> 'b * 's) -> 'b t * 's
 val fold2_exn : 'a t -> 'b t -> 's -> f:('a -> 'b -> 's -> 's) -> 's