index.cpp

   1 #include "stdafx.h"
   2 #include "defs.h"
   3
   4 // n is the total number of things on the stack. The first thing on the stack
   5 // is the object to be indexed, followed by the indices themselves.
   6
   7 void
   8 index_function(int n)
   9 {
  10         int i, k, m, ndim, nelem, t;
  11         U **s;
  12
  13         save();
  14         s = stack + tos - n;
  15         p1 = s[0];
  16
  17         // index of scalar ok
  18
  19         if (!istensor(p1)) {
  20                 tos -= n;
  21                 push(p1);
  22                 restore();
  23                 return;
  24         }
  25
  26         ndim = p1->u.tensor->ndim;
  27
  28         m = n - 1;
  29
  30         if (m > ndim)
  31                 stop("too many indices for tensor");
  32
  33         k = 0;
  34
  35         for (i = 0; i < m; i++) {
  36                 push(s[i + 1]);
  37                 t = pop_integer();
  38                 if (t < 1 || t > p1->u.tensor->dim[i])
  39                         stop("index out of range");
  40                 k = k * p1->u.tensor->dim[i] + t - 1;
  41         }
  42
  43         if (ndim == m) {
  44                 tos -= n;
  45                 push(p1->u.tensor->elem[k]);
  46                 restore();
  47                 return;
  48         }
  49
  50         for (i = m; i < ndim; i++)
  51                 k = k * p1->u.tensor->dim[i] + 0;
  52
  53         nelem = 1;
  54
  55         for (i = m; i < ndim; i++)
  56                 nelem *= p1->u.tensor->dim[i];
  57
  58         p2 = alloc_tensor(nelem);
  59
  60         p2->u.tensor->ndim = ndim - m;
  61
  62         for (i = m; i < ndim; i++)
  63                 p2->u.tensor->dim[i - m] = p1->u.tensor->dim[i];
  64
  65         for (i = 0; i < nelem; i++)
  66                 p2->u.tensor->elem[i] = p1->u.tensor->elem[k + i];
  67
  68         tos -= n;
  69         push(p2);
  70         restore();
  71 }
  72
  73 //-----------------------------------------------------------------------------
  74 //
  75 //      Input:          n               Number of args on stack
  76 //
  77 //                      tos-n           Right-hand value
  78 //
  79 //                      tos-n+1         Left-hand value
  80 //
  81 //                      tos-n+2         First index
  82 //
  83 //                      .
  84 //                      .
  85 //                      .
  86 //
  87 //                      tos-1           Last index
  88 //
  89 //      Output:         Result on stack
  90 //
  91 //-----------------------------------------------------------------------------
  92
  93 #define LVALUE p1
  94 #define RVALUE p2
  95 #define TMP p3
  96
  97 void
  98 set_component(int n)
  99 {
 100         int i, k, m, ndim, t;
 101         U **s;
 102
 103         save();
 104
 105         if (n < 3)
 106                 stop("error in indexed assign");
 107
 108         s = stack + tos - n;
 109
 110         RVALUE = s[0];
 111
 112         LVALUE = s[1];
 113
 114         if (!istensor(LVALUE))
 115                 stop("error in indexed assign");
 116
 117         ndim = LVALUE->u.tensor->ndim;
 118
 119         m = n - 2;
 120
 121         if (m > ndim)
 122                 stop("error in indexed assign");
 123
 124         k = 0;
 125
 126         for (i = 0; i < m; i++) {
 127                 push(s[i + 2]);
 128                 t = pop_integer();
 129                 if (t < 1 || t > LVALUE->u.tensor->dim[i])
 130                         stop("error in indexed assign");
 131                 k = k * p1->u.tensor->dim[i] + t - 1;
 132         }
 133
 134         for (i = m; i < ndim; i++)
 135                 k = k * p1->u.tensor->dim[i] + 0;
 136
 137         // copy
 138
 139         TMP = alloc_tensor(LVALUE->u.tensor->nelem);
 140
 141         TMP->u.tensor->ndim = LVALUE->u.tensor->ndim;
 142
 143         for (i = 0; i < p1->u.tensor->ndim; i++)
 144                 TMP->u.tensor->dim[i] = LVALUE->u.tensor->dim[i];
 145
 146         for (i = 0; i < p1->u.tensor->nelem; i++)
 147                 TMP->u.tensor->elem[i] = LVALUE->u.tensor->elem[i];
 148
 149         LVALUE = TMP;
 150
 151         if (ndim == m) {
 152                 if (istensor(RVALUE))
 153                         stop("error in indexed assign");
 154                 LVALUE->u.tensor->elem[k] = RVALUE;
 155                 tos -= n;
 156                 push(LVALUE);
 157                 restore();
 158                 return;
 159         }
 160
 161         // see if the rvalue matches
 162
 163         if (!istensor(RVALUE))
 164                 stop("error in indexed assign");
 165
 166         if (ndim - m != RVALUE->u.tensor->ndim)
 167                 stop("error in indexed assign");
 168
 169         for (i = 0; i < RVALUE->u.tensor->ndim; i++)
 170                 if (LVALUE->u.tensor->dim[m + i] != RVALUE->u.tensor->dim[i])
 171                         stop("error in indexed assign");
 172
 173         // copy rvalue
 174
 175         for (i = 0; i < RVALUE->u.tensor->nelem; i++)
 176                 LVALUE->u.tensor->elem[k + i] = RVALUE->u.tensor->elem[i];
 177
 178         tos -= n;
 179
 180         push(LVALUE);
 181
 182         restore();
 183 }
 184
 185 #if SELFTEST
 186
 187 static char *s[] = {
 188
 189         "A=((A11,A12),(A21,A22))",
 190         "",
 191
 192         "A[1,1]",
 193         "A11",
 194
 195         "A[1,2]",
 196         "A12",
 197
 198         "A[2,1]",
 199         "A21",
 200
 201         "A[2,2]",
 202         "A22",
 203
 204         "A[1]",
 205         "(A11,A12)",
 206
 207         "A[2]",
 208         "(A21,A22)",
 209
 210         "A[1]=(B11,B12)",
 211         "",
 212
 213         "A",
 214         "((B11,B12),(A21,A22))",
 215
 216         "A[2]=(B21,B22)",
 217         "",
 218
 219         "A",
 220         "((B11,B12),(B21,B22))",
 221
 222         "A=((0,0),(0,0))",
 223         "",
 224
 225         "A[1,1]",
 226         "0",
 227
 228         // index of scalar ok
 229
 230         "1[2]",
 231         "1",
 232 };
 233
 234 void
 235 test_index(void)
 236 {
 237         test(__FILE__, s, sizeof s / sizeof (char *));
 238 }
 239
 240 #endif