sys/src/cmd/fplot.c

   1 #include <u.h>
   2 #include <libc.h>
   3 #include <ctype.h>
   4 #include <draw.h>
   5 #include <event.h>
   6
   7 typedef struct Op Op;
   8 typedef struct Operator Operator;
   9 typedef struct Token Token;
  10 typedef struct Constant Constant;
  11 typedef struct Code Code;
  12
  13 enum {
  14         ONONE,
  15         ONUMBER,
  16         OVAR,
  17         OUNARY,
  18         OBINARY,
  19 };
  20
  21 struct Op {
  22         int type;
  23         union {
  24                 void (*f)(void);
  25                 double val;
  26         };
  27 };
  28
  29 enum {
  30         TNONE,
  31         TNUMBER,
  32         TVAR,
  33         TOP,
  34         TPARENL,
  35         TPARENR,
  36 };
  37
  38 struct Token {
  39         int type;
  40         union {
  41                 Operator *op;
  42                 double val;
  43         };
  44         Token *next;
  45 };
  46
  47 double *stack, *sp;
  48 void add(void) { sp--; *sp += *(sp+1); }
  49 void sub(void) { sp--; *sp -= *(sp+1); }
  50 void mul(void) { sp--; *sp *= *(sp+1); }
  51 void div(void) { sp--; *sp /= *(sp+1); }
  52 void pot(void) { sp--; *sp = pow(*sp, *(sp+1)); }
  53 void osin(void) { *sp = sin(*sp); }
  54 void ocos(void) { *sp = cos(*sp); }
  55 void otan(void) { *sp = tan(*sp); }
  56 void oasin(void) { *sp = asin(*sp); }
  57 void oacos(void) { *sp = acos(*sp); }
  58 void oatan(void) { *sp = atan(*sp); }
  59 void osqrt(void) { *sp = sqrt(*sp); }
  60 void olog(void) { *sp = log10(*sp); }
  61 void oln(void) { *sp = log(*sp); }
  62
  63 struct Operator {
  64         char *s;
  65         char type;
  66         char rassoc;
  67         short prec;
  68         void (*f)(void);
  69 } ops[] = {
  70         "+",    OBINARY,        0,      0,      add,
  71         "-",    OBINARY,        0,      0,      sub,
  72         "*",    OBINARY,        0,      100,    mul,
  73         "/",    OBINARY,        0,      100,    div,
  74         "^",    OBINARY,        1,      200,    pot,
  75         "sin",  OUNARY,         0,      50,     osin,
  76         "cos",  OUNARY,         0,      50,     ocos,
  77         "tan",  OUNARY,         0,      50,     otan,
  78         "asin", OUNARY,         0,      50,     oasin,
  79         "acos", OUNARY,         0,      50,     oacos,
  80         "atan", OUNARY,         0,      50,     oatan,
  81         "sqrt", OUNARY,         0,      50,     osqrt,
  82         "log",  OUNARY,         0,      50,     olog,
  83         "ln",   OUNARY,         0,      50,     oln,
  84 };
  85
  86 struct Constant {
  87         char *s;
  88         double val;
  89 } consts[] = {
  90         "pi",   3.14159265359,
  91         "π",   3.14159265359,
  92         "e",    2.71828182846,
  93 };
  94
  95 struct Code {
  96         Op* p;
  97         int len, cap;
  98 } *fns;
  99 int nfns;
 100
 101 Token *opstackbot;
 102 double xmin = -10, xmax = 10;
 103 double ymin = -10, ymax = 10;
 104 Image *color;
 105 int cflag;
 106 char *imagedata;
 107 int picx = 640, picy = 480;
 108
 109 void *
 110 emalloc(int size)
 111 {
 112         void *v;
 113
 114         v = mallocz(size, 1);
 115         if(v == nil)
 116                 sysfatal("emalloc: %r");
 117         setmalloctag(v, getcallerpc(&size));
 118         return v;
 119 }
 120
 121 void
 122 addop(Code *c, Op o)
 123 {
 124         if(c->len >= c->cap) {
 125                 c->cap += 32;
 126                 c->p = realloc(c->p, sizeof(Op) * c->cap);
 127                 if(c->p == nil)
 128                         sysfatal("realloc: %r");
 129         }
 130         c->p[c->len++] = o;
 131 }
 132
 133 void
 134 push(Token *t)
 135 {
 136         t->next = opstackbot;
 137         opstackbot = t;
 138 }
 139
 140 void
 141 pop(Code *c)
 142 {
 143         Token *t;
 144         Op o;
 145
 146         t = opstackbot;
 147         if(t == nil)
 148                 sysfatal("stack underflow");
 149         opstackbot = t->next;
 150         if(t->type != TOP)
 151                 sysfatal("non-operator pop");
 152         o.type = t->op->type;
 153         o.f = t->op->f;
 154         addop(c, o);
 155         free(t);
 156 }
 157
 158 void
 159 popdel(void)
 160 {
 161         Token *t;
 162
 163         t = opstackbot;
 164         opstackbot = t->next;
 165         free(t);
 166 }
 167
 168 Token *
 169 lex(char **s)
 170 {
 171         Token *t;
 172         Operator *o;
 173         Constant *c;
 174
 175         while(isspace(**s))
 176                 (*s)++;
 177         if(**s == 0)
 178                 return nil;
 179
 180         t = emalloc(sizeof(*t));
 181         if(isdigit(**s)) {
 182                 t->type = TNUMBER;
 183                 t->val = strtod(*s, s);
 184                 return t;
 185         }
 186         if(**s == '(') {
 187                 t->type = TPARENL;
 188                 (*s)++;
 189                 return t;
 190         }
 191         if(**s == ')') {
 192                 t->type = TPARENR;
 193                 (*s)++;
 194                 return t;
 195         }
 196         if(**s == 'x') {
 197                 t->type = TVAR;
 198                 (*s)++;
 199                 return t;
 200         }
 201         for(o = ops; o < ops + nelem(ops); o++)
 202                 if(strncmp(*s, o->s, strlen(o->s)) == 0) {
 203                         t->type = TOP;
 204                         t->op = o;
 205                         *s += strlen(o->s);
 206                         return t;
 207                 }
 208         for(c = consts; c < consts + nelem(consts); c++)
 209                 if(strncmp(*s, c->s, strlen(c->s)) == 0) {
 210                         t->type = TNUMBER;
 211                         t->val = c->val;
 212                         *s += strlen(c->s);
 213                         return t;
 214                 }
 215         sysfatal("syntax error at %s", *s);
 216         return nil;
 217 }
 218
 219 void
 220 parse(Code *c, char *s)
 221 {
 222         Token *t;
 223         Op o;
 224
 225         while(t = lex(&s)) {
 226                 switch(t->type) {
 227                 case TNUMBER:
 228                         o.type = ONUMBER;
 229                         o.val = t->val;
 230                         addop(c, o);
 231                         free(t);
 232                         break;
 233                 case TVAR:
 234                         o.type = OVAR;
 235                         addop(c, o);
 236                         free(t);
 237                         break;
 238                 case TOP:
 239                         if(t->op->type == OBINARY)
 240                                 while(opstackbot != nil && opstackbot->type == TOP &&
 241                                         (opstackbot->op->prec > t->op->prec ||
 242                                         t->op->rassoc && opstackbot->op->prec == t->op->prec))
 243                                         pop(c);
 244                         push(t);
 245                         break;
 246                 case TPARENL:
 247                         push(t);
 248                         break;
 249                 case TPARENR:
 250                         while(opstackbot != nil && opstackbot->type == TOP)
 251                                 pop(c);
 252                         if(opstackbot == nil)
 253                                 sysfatal("mismatched parentheses");
 254                         popdel();
 255                         free(t);
 256                         break;
 257                 default:
 258                         sysfatal("unknown token type %d", t->type);
 259                 }
 260         }
 261         while(opstackbot != nil)
 262                 switch(opstackbot->type) {
 263                 case TOP:
 264                         pop(c);
 265                         break;
 266                 case TPARENL:
 267                         sysfatal("mismatched parentheses");
 268                 default:
 269                         sysfatal("syntax error");
 270                 }
 271 }
 272
 273 int
 274 calcstack(Code *c)
 275 {
 276         int cur, max;
 277         Op *o;
 278
 279         cur = 0;
 280         max = 0;
 281         for(o = c->p; o < c->p + c->len; o++)
 282                 switch(o->type) {
 283                 case ONUMBER: case OVAR:
 284                         if(++cur > max)
 285                                 max = cur;
 286                         break;
 287                 case OUNARY:
 288                         if(cur == 0)
 289                                 sysfatal("syntax error");
 290                         break;
 291                 case OBINARY:
 292                         if(cur <= 1)
 293                                 sysfatal("syntax error");
 294                         cur--;
 295                         break;
 296                 }
 297
 298         if(cur != 1)
 299                 sysfatal("syntax error");
 300         return max;
 301 }
 302
 303 double
 304 calc(Code *c, double x)
 305 {
 306         Op *o;
 307
 308         sp = stack - 1;
 309         for(o = c->p; o < c->p + c->len; o++)
 310                 switch(o->type) {
 311                 case ONUMBER:
 312                         *++sp = o->val;
 313                         break;
 314                 case OVAR:
 315                         *++sp = x;
 316                         break;
 317                 case OUNARY: case OBINARY:
 318                         o->f();
 319                 }
 320         return *sp;
 321 }
 322
 323 double
 324 convx(Rectangle *r, int x)
 325 {
 326         return (xmax - xmin) * (x - r->min.x) / (r->max.x - r->min.x) + xmin;
 327 }
 328
 329 int
 330 deconvx(Rectangle *r, double dx)
 331 {
 332         return (dx - xmin) * (r->max.x - r->min.x) / (xmax - xmin) + r->min.x + 0.5;
 333 }
 334
 335 double
 336 convy(Rectangle *r, int y)
 337 {
 338         return (ymax - ymin) * (r->max.y - y) / (r->max.y - r->min.y) + ymin;
 339 }
 340
 341 int
 342 deconvy(Rectangle *r, double dy)
 343 {
 344         return (ymax - dy) * (r->max.y - r->min.y) / (ymax - ymin) + r->min.y + 0.5;
 345 }
 346
 347 void
 348 pixel(int x, int y)
 349 {
 350         char *p;
 351
 352         if(cflag) {
 353                 if(x >= picx || y >= picy || x < 0 || y < 0)
 354                         return;
 355                 p = imagedata + (picx * y + x) * 3;
 356                 p[0] = p[1] = p[2] = 0;
 357         } else
 358                 draw(screen, Rect(x, y, x + 1, y + 1), color, nil, ZP);
 359 }
 360
 361 void
 362 drawinter(Code *co, Rectangle *r, double x1, double x2, int n)
 363 {
 364         double y1, y2;
 365         int iy1, iy2;
 366         int ix1, ix2;
 367
 368         ix1 = deconvx(r, x1);
 369         ix2 = deconvx(r, x2);
 370         iy1 = iy2 = 0;
 371         y1 = calc(co, x1);
 372         if(!isNaN(y1)) {
 373                 iy1 = deconvy(r, y1);
 374                 pixel(ix1, iy1);
 375         }
 376         y2 = calc(co, x2);
 377         if(!isNaN(y2)) {
 378                 iy2 = deconvy(r, y2);
 379                 pixel(ix2, iy2);
 380         }
 381         if(isNaN(y1) || isNaN(y2))
 382                 return;
 383         if(n >= 10)
 384                 return;
 385         if(iy2 >= iy1 - 1 && iy2 <= iy1 + 1)
 386                 return;
 387         if(iy1 > r->max.y && iy2 > r->max.y)
 388                 return;
 389         if(iy1 < r->min.y && iy2 < r->min.y)
 390                 return;
 391         drawinter(co, r, x1, (x1 + x2) / 2, n + 1);
 392         drawinter(co, r, (x1 + x2) / 2, x2, n + 1);
 393 }
 394
 395 void
 396 drawgraph(Code *co, Rectangle *r)
 397 {
 398         int x;
 399
 400         for(x = r->min.x; x < r->max.x; x++)
 401                 drawinter(co, r, convx(r, x), convx(r, x + 1), 0);
 402 }
 403
 404 void
 405 drawgraphs(void)
 406 {
 407         int i;
 408
 409         color = display->black;
 410         for(i = 0; i < nfns; i++)
 411                 drawgraph(&fns[i], &screen->r);
 412         flushimage(display, 1);
 413 }
 414
 415 void
 416 usage(void)
 417 {
 418         fprint(2, "usage: fplot [-c [-s size]] [-r range] functions ...\n");
 419         exits("usage");
 420 }
 421
 422 void
 423 zoom(void)
 424 {
 425         Mouse m;
 426         Rectangle r;
 427         double xmin_, xmax_, ymin_, ymax_;
 428
 429         m.buttons = 7;
 430         r = egetrect(1, &m);
 431         if(r.min.x == 0 && r.min.y == 0 && r.max.x == 0 && r.max.y == 0)
 432                 return;
 433         xmin_ = convx(&screen->r, r.min.x);
 434         xmax_ = convx(&screen->r, r.max.x);
 435         ymin_ = convy(&screen->r, r.max.y);
 436         ymax_ = convy(&screen->r, r.min.y);
 437         xmin = xmin_;
 438         xmax = xmax_;
 439         ymin = ymin_;
 440         ymax = ymax_;
 441         draw(screen, screen->r, display->white, nil, ZP);
 442         drawgraphs();
 443 }
 444
 445 void
 446 parsefns(int n, char **s)
 447 {
 448         int i, max, cur;
 449
 450         max = 0;
 451         nfns = n;
 452         fns = emalloc(sizeof(*fns) * n);
 453         for(i = 0; i < nfns; i++) {
 454                 parse(&fns[i], s[i]);
 455                 cur = calcstack(&fns[i]);
 456                 if(cur > max)
 457                         max = cur;
 458         }
 459         stack = emalloc(sizeof(*stack) * max);
 460 }
 461
 462 void
 463 parserange(char *s)
 464 {
 465         while(*s && !isdigit(*s)) s++;
 466         if(*s == 0) return;
 467         xmin = strtod(s, &s);
 468         while(*s && !isdigit(*s)) s++;
 469         if(*s == 0) return;
 470         xmax = strtod(s, &s);
 471         while(*s && !isdigit(*s)) s++;
 472         if(*s == 0) return;
 473         ymin = strtod(s, &s);
 474         while(*s && !isdigit(*s)) s++;
 475         if(*s == 0) return;
 476         ymax = strtod(s, &s);
 477 }
 478
 479 void
 480 parsesize(char *s)
 481 {
 482         while(*s && !isdigit(*s)) s++;
 483         if(*s == 0) return;
 484         picx = strtol(s, &s, 0);
 485         while(*s && !isdigit(*s)) s++;
 486         if(*s == 0) return;
 487         picy = strtol(s, &s, 0);
 488 }
 489
 490 void
 491 main(int argc, char **argv)
 492 {
 493         Event e;
 494         Rectangle r;
 495         int i;
 496
 497         ARGBEGIN {
 498         case 'r': parserange(EARGF(usage())); break;
 499         case 's': parsesize(EARGF(usage())); break;
 500         case 'c': cflag++; break;
 501         default: usage();
 502         } ARGEND;
 503         if(argc < 1)
 504                 usage();
 505         setfcr(getfcr() & ~(FPZDIV | FPINVAL));
 506         parsefns(argc, argv);
 507         if(cflag) {
 508                 imagedata = emalloc(picx * picy * 3);
 509                 memset(imagedata, 0xFF, picx * picy * 3);
 510                 print("%11s %11d %11d %11d %11d ", "r8g8b8", 0, 0, picx, picy);
 511                 r.min.x = r.min.y = 0;
 512                 r.max.x = picx;
 513                 r.max.y = picy;
 514                 for(i = 0; i < nfns; i++)
 515                         drawgraph(&fns[i], &r);
 516                 if(write(1, imagedata, picx * picy * 3) < picx * picy * 3)
 517                         sysfatal("write: %r");
 518         } else {
 519                 if(initdraw(nil, nil, "fplot") < 0)
 520                         sysfatal("initdraw: %r");
 521                 einit(Emouse | Ekeyboard);
 522                 drawgraphs();
 523                 for(;;) {
 524                         switch(event(&e)) {
 525                         case Ekeyboard:
 526                                 switch(e.kbdc) {
 527                                 case 'q': exits(nil);
 528                                 case 'z': zoom();
 529                                 }
 530                         }
 531                 }
 532         }
 533 }
 534
 535 void
 536 eresized(int new)
 537 {
 538         if(new) {
 539                 if(getwindow(display, Refnone) < 0)
 540                         sysfatal("getwindow: %r");
 541                 drawgraphs();
 542         }
 543 }