From 93a147af7b5fda7f94f135ba6d97bf9d7ad7ee23 Mon Sep 17 00:00:00 2001
From: peter <petrel.harp@gmail.com>
Date: Tue, 24 Sep 2024 17:18:55 -0700
Subject: [PATCH] more restricts

---
 c/tskit/trees.c | 36 +++++++++++++++++++++---------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/c/tskit/trees.c b/c/tskit/trees.c
index bc1db29c7b..9462dfc51c 100644
--- a/c/tskit/trees.c
+++ b/c/tskit/trees.c
@@ -10108,10 +10108,16 @@ tsk_matvec_calculator_write_output(tsk_matvec_calculator_t *self, double *restri
     tsk_id_t u;
     tsk_size_t j, k;
     tsk_size_t n = tsk_treeseq_get_num_samples(self->ts);
+    const tsk_size_t num_weights = self->num_weights;
+    const double tree_left = self->tree_left;
     double *u_row, *out_row;
-    double *out_means = tsk_malloc(self->num_weights * sizeof(*out_means));
+    double *out_means = tsk_malloc(num_weights * sizeof(*out_means));
     const tsk_id_t *restrict parent = self->parent;
     const double *restrict nodes_time = self->ts->tables->nodes.time;
+    double *restrict x = self->x;
+    double *restrict w = self->w;
+    double *restrict v = self->v;
+    const tsk_id_t *restrict samples = self->ts->samples;
 
     if (out_means == NULL) {
         ret = TSK_ERR_NO_MEMORY;
@@ -10119,15 +10125,15 @@ tsk_matvec_calculator_write_output(tsk_matvec_calculator_t *self, double *restri
     }
 
     for (j = 0; j < n; j++) {
-        out_row = GET_2D_ROW(y, self->num_weights, j);
-        u = self->ts->samples[j];
+        out_row = GET_2D_ROW(y, num_weights, j);
+        u = samples[j];
         while (u != TSK_NULL) {
-            if (self->x[u] != self->tree_left) {
-                tsk_matvec_calculator_add_z(u, parent[u], self->tree_left, self->x,
-                    self->num_weights, self->w, self->v, nodes_time);
+            if (x[u] != tree_left) {
+                tsk_matvec_calculator_add_z(
+                    u, parent[u], tree_left, x, num_weights, w, v, nodes_time);
             }
-            u_row = GET_2D_ROW(self->v, self->num_weights, u);
-            for (k = 0; k < self->num_weights; k++) {
+            u_row = GET_2D_ROW(v, num_weights, u);
+            for (k = 0; k < num_weights; k++) {
                 out_row[k] += u_row[k];
             }
             u = parent[u];
@@ -10135,27 +10141,27 @@ tsk_matvec_calculator_write_output(tsk_matvec_calculator_t *self, double *restri
     }
 
     if (!(self->options & TSK_STAT_NONCENTRED)) {
-        for (k = 0; k < self->num_weights; k++) {
+        for (k = 0; k < num_weights; k++) {
             out_means[k] = 0.0;
         }
         for (j = 0; j < n; j++) {
-            out_row = GET_2D_ROW(y, self->num_weights, j);
-            for (k = 0; k < self->num_weights; k++) {
+            out_row = GET_2D_ROW(y, num_weights, j);
+            for (k = 0; k < num_weights; k++) {
                 out_means[k] += out_row[k];
             }
         }
-        for (k = 0; k < self->num_weights; k++) {
+        for (k = 0; k < num_weights; k++) {
             out_means[k] /= (double) n;
         }
         for (j = 0; j < n; j++) {
-            out_row = GET_2D_ROW(y, self->num_weights, j);
-            for (k = 0; k < self->num_weights; k++) {
+            out_row = GET_2D_ROW(y, num_weights, j);
+            for (k = 0; k < num_weights; k++) {
                 out_row[k] -= out_means[k];
             }
         }
     }
     /* zero out v */
-    tsk_memset(self->v, 0, self->num_nodes * self->num_weights * sizeof(*self->v));
+    tsk_memset(self->v, 0, self->num_nodes * num_weights * sizeof(*self->v));
 out:
     tsk_safe_free(out_means);
     return ret;