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 mod(void) { sp--; *sp = fmod(*sp, *(sp+1)); }
  53 void pot(void) { sp--; *sp = pow(*sp, *(sp+1)); }
  54 void osin(void) { *sp = sin(*sp); }
  55 void ocos(void) { *sp = cos(*sp); }
  56 void otan(void) { *sp = tan(*sp); }
  57 void oasin(void) { *sp = asin(*sp); }
  58 void oacos(void) { *sp = acos(*sp); }
  59 void oatan(void) { *sp = atan(*sp); }
  60 void osqrt(void) { *sp = sqrt(*sp); }
  61 void oexp(void) { *sp = exp(*sp); }
  62 void olog(void) { *sp = log10(*sp); }
  63 void oln(void) { *sp = log(*sp); }
  64
  65 struct Operator {
  66         char *s;
  67         char type;
  68         char rassoc;
  69         short prec;
  70         void (*f)(void);
  71 } ops[] = {
  72         "+",    OBINARY,        0,      0,      add,
  73         "-",    OBINARY,        0,      0,      sub,
  74         "*",    OBINARY,        0,      100,    mul,
  75         "/",    OBINARY,        0,      100,    div,
  76         "%",    OBINARY,        0,      100,    mod,
  77         "^",    OBINARY,        1,      200,    pot,
  78         "sin",  OUNARY,         0,      300,    osin,
  79         "cos",  OUNARY,         0,      300,    ocos,
  80         "tan",  OUNARY,         0,      300,    otan,
  81         "asin", OUNARY,         0,      300,    oasin,
  82         "acos", OUNARY,         0,      300,    oacos,
  83         "atan", OUNARY,         0,      300,    oatan,
  84         "sqrt", OUNARY,         0,      300,    osqrt,
  85         "exp",  OUNARY,         0,      300,    oexp,
  86         "log",  OUNARY,         0,      300,    olog,
  87         "ln",   OUNARY,         0,      300,    oln,
  88 };
  89
  90 struct Constant {
  91         char *s;
  92         double val;
  93 } consts[] = {
  94         "pi",   3.14159265359,
  95         "π",   3.14159265359,
  96         "e",    2.71828182846,
  97 };
  98
  99 struct Code {
 100         Op* p;
 101         int len, cap;
 102 } *fns;
 103 int nfns;
 104
 105 Token *opstackbot;
 106 double xmin = -10, xmax = 10;
 107 double ymin = -10, ymax = 10;
 108 double gymin, gymax;
 109 Image *color;
 110 int cflag, aflag;
 111 char *imagedata;
 112 int picx = 640, picy = 480;
 113
 114 typedef struct FRectangle FRectangle;
 115 struct FRectangle {
 116         double xmin, xmax, ymin, ymax;
 117 } *zoomst;
 118 int nzoomst;
 119
 120 void *
 121 emalloc(int size)
 122 {
 123         void *v;
 124
 125         v = mallocz(size, 1);
 126         if(v == nil)
 127                 sysfatal("emalloc: %r");
 128         setmalloctag(v, getcallerpc(&size));
 129         return v;
 130 }
 131
 132 void
 133 addop(Code *c, Op o)
 134 {
 135         if(c->len >= c->cap) {
 136                 c->cap += 32;
 137                 c->p = realloc(c->p, sizeof(Op) * c->cap);
 138                 if(c->p == nil)
 139                         sysfatal("realloc: %r");
 140         }
 141         c->p[c->len++] = o;
 142 }
 143
 144 void
 145 push(Token *t)
 146 {
 147         t->next = opstackbot;
 148         opstackbot = t;
 149 }
 150
 151 void
 152 pop(Code *c)
 153 {
 154         Token *t;
 155         Op o;
 156
 157         t = opstackbot;
 158         if(t == nil)
 159                 sysfatal("stack underflow");
 160         opstackbot = t->next;
 161         if(t->type != TOP)
 162                 sysfatal("non-operator pop");
 163         o.type = t->op->type;
 164         o.f = t->op->f;
 165         addop(c, o);
 166         free(t);
 167 }
 168
 169 void
 170 popdel(void)
 171 {
 172         Token *t;
 173
 174         t = opstackbot;
 175         opstackbot = t->next;
 176         free(t);
 177 }
 178
 179 Token *
 180 lex(char **s)
 181 {
 182         Token *t;
 183         Operator *o;
 184         Constant *c;
 185
 186         while(isspace(**s))
 187                 (*s)++;
 188         if(**s == 0)
 189                 return nil;
 190
 191         t = emalloc(sizeof(*t));
 192         if(isdigit(**s)) {
 193                 t->type = TNUMBER;
 194                 t->val = strtod(*s, s);
 195                 return t;
 196         }
 197         if(**s == '(') {
 198                 t->type = TPARENL;
 199                 (*s)++;
 200                 return t;
 201         }
 202         if(**s == ')') {
 203                 t->type = TPARENR;
 204                 (*s)++;
 205                 return t;
 206         }
 207         if(**s == 'x') {
 208                 t->type = TVAR;
 209                 (*s)++;
 210                 return t;
 211         }
 212         for(o = ops; o < ops + nelem(ops); o++)
 213                 if(strncmp(*s, o->s, strlen(o->s)) == 0) {
 214                         t->type = TOP;
 215                         t->op = o;
 216                         *s += strlen(o->s);
 217                         return t;
 218                 }
 219         for(c = consts; c < consts + nelem(consts); c++)
 220                 if(strncmp(*s, c->s, strlen(c->s)) == 0) {
 221                         t->type = TNUMBER;
 222                         t->val = c->val;
 223                         *s += strlen(c->s);
 224                         return t;
 225                 }
 226         sysfatal("syntax error at %s", *s);
 227         return nil;
 228 }
 229
 230 void
 231 parse(Code *c, char *s)
 232 {
 233         Token *t;
 234         Op o;
 235
 236         while(t = lex(&s)) {
 237                 switch(t->type) {
 238                 case TNUMBER:
 239                         o.type = ONUMBER;
 240                         o.val = t->val;
 241                         addop(c, o);
 242                         free(t);
 243                         break;
 244                 case TVAR:
 245                         o.type = OVAR;
 246                         addop(c, o);
 247                         free(t);
 248                         break;
 249                 case TOP:
 250                         if(t->op->type == OBINARY)
 251                                 while(opstackbot != nil && opstackbot->type == TOP &&
 252                                         (opstackbot->op->prec > t->op->prec ||
 253                                         t->op->rassoc && opstackbot->op->prec == t->op->prec))
 254                                         pop(c);
 255                         push(t);
 256                         break;
 257                 case TPARENL:
 258                         push(t);
 259                         break;
 260                 case TPARENR:
 261                         while(opstackbot != nil && opstackbot->type == TOP)
 262                                 pop(c);
 263                         if(opstackbot == nil)
 264                                 sysfatal("mismatched parentheses");
 265                         popdel();
 266                         free(t);
 267                         break;
 268                 default:
 269                         sysfatal("unknown token type %d", t->type);
 270                 }
 271         }
 272         while(opstackbot != nil)
 273                 switch(opstackbot->type) {
 274                 case TOP:
 275                         pop(c);
 276                         break;
 277                 case TPARENL:
 278                         sysfatal("mismatched parentheses");
 279                 default:
 280                         sysfatal("syntax error");
 281                 }
 282 }
 283
 284 int
 285 calcstack(Code *c)
 286 {
 287         int cur, max;
 288         Op *o;
 289
 290         cur = 0;
 291         max = 0;
 292         for(o = c->p; o < c->p + c->len; o++)
 293                 switch(o->type) {
 294                 case ONUMBER: case OVAR:
 295                         if(++cur > max)
 296                                 max = cur;
 297                         break;
 298                 case OUNARY:
 299                         if(cur == 0)
 300                                 sysfatal("syntax error");
 301                         break;
 302                 case OBINARY:
 303                         if(cur <= 1)
 304                                 sysfatal("syntax error");
 305                         cur--;
 306                         break;
 307                 }
 308
 309         if(cur != 1)
 310                 sysfatal("syntax error");
 311         return max;
 312 }
 313
 314 double
 315 calc(Code *c, double x)
 316 {
 317         Op *o;
 318
 319         sp = stack - 1;
 320         for(o = c->p; o < c->p + c->len; o++)
 321                 switch(o->type) {
 322                 case ONUMBER:
 323                         *++sp = o->val;
 324                         break;
 325                 case OVAR:
 326                         *++sp = x;
 327                         break;
 328                 case OUNARY: case OBINARY:
 329                         o->f();
 330                 }
 331         if(*sp < gymin) gymin = *sp;
 332         if(*sp > gymax) gymax = *sp;
 333         return *sp;
 334 }
 335
 336 double
 337 convx(Rectangle *r, int x)
 338 {
 339         return (xmax - xmin) * (x - r->min.x) / (r->max.x - r->min.x) + xmin;
 340 }
 341
 342 int
 343 deconvx(Rectangle *r, double dx)
 344 {
 345         return (dx - xmin) * (r->max.x - r->min.x) / (xmax - xmin) + r->min.x + 0.5;
 346 }
 347
 348 double
 349 convy(Rectangle *r, int y)
 350 {
 351         return (ymax - ymin) * (r->max.y - y) / (r->max.y - r->min.y) + ymin;
 352 }
 353
 354 int
 355 deconvy(Rectangle *r, double dy)
 356 {
 357         return (ymax - dy) * (r->max.y - r->min.y) / (ymax - ymin) + r->min.y + 0.5;
 358 }
 359
 360 void
 361 pixel(int x, int y)
 362 {
 363         char *p;
 364
 365         if(cflag) {
 366                 if(x >= picx || y >= picy || x < 0 || y < 0)
 367                         return;
 368                 p = imagedata + (picx * y + x) * 3;
 369                 p[0] = p[1] = p[2] = 0;
 370         } else
 371                 draw(screen, Rect(x, y, x + 1, y + 1), color, nil, ZP);
 372 }
 373
 374 void
 375 drawinter(Code *co, Rectangle *r, double x1, double x2, int n)
 376 {
 377         double y1, y2;
 378         int iy1, iy2;
 379         int ix1, ix2;
 380
 381         ix1 = deconvx(r, x1);
 382         ix2 = deconvx(r, x2);
 383         iy1 = iy2 = 0;
 384         y1 = calc(co, x1);
 385         if(!isNaN(y1)) {
 386                 iy1 = deconvy(r, y1);
 387                 pixel(ix1, iy1);
 388         }
 389         y2 = calc(co, x2);
 390         if(!isNaN(y2)) {
 391                 iy2 = deconvy(r, y2);
 392                 pixel(ix2, iy2);
 393         }
 394         if(isNaN(y1) || isNaN(y2))
 395                 return;
 396         if(n >= 10)
 397                 return;
 398         if(iy2 >= iy1 - 1 && iy2 <= iy1 + 1)
 399                 return;
 400         if(iy1 > r->max.y && iy2 > r->max.y)
 401                 return;
 402         if(iy1 < r->min.y && iy2 < r->min.y)
 403                 return;
 404         drawinter(co, r, x1, (x1 + x2) / 2, n + 1);
 405         drawinter(co, r, (x1 + x2) / 2, x2, n + 1);
 406 }
 407
 408 void
 409 drawgraph(Code *co, Rectangle *r)
 410 {
 411         int x;
 412
 413         gymin = Inf(1);
 414         gymax = Inf(-1);
 415         for(x = r->min.x; x < r->max.x; x++)
 416                 drawinter(co, r, convx(r, x), convx(r, x + 1), 0);
 417 }
 418
 419 void
 420 tickfmt(double d, double m, int n, char *fmt)
 421 {
 422         double e1, e2;
 423         int x;
 424
 425         e1 = log10(fabs(m));
 426         e2 = log10(fabs(m + n * d));
 427         if(e2 > e1) e1 = e2;
 428         if(e2 >= 4 || e2 < -3){
 429                 x = ceil(e1-log10(d)-1);
 430                 snprint(fmt, 32, "%%.%de", x);
 431         }else{
 432                 x = ceil(-log10(d));
 433                 snprint(fmt, 32, "%%.%df", x);
 434         }
 435 }
 436
 437 int
 438 xticklabel(char *fmt, double g, int p, int x, int y)
 439 {
 440         char buf[32];
 441         Rectangle lr;
 442         int ny;
 443
 444         snprint(buf, sizeof(buf), fmt, g);
 445         lr.min = Pt(p, y+2);
 446         lr.max = addpt(lr.min, stringsize(display->defaultfont, buf));
 447         lr = rectsubpt(lr, Pt(Dx(lr) / 2-1, 0));
 448         if(lr.max.y >= screen->r.max.y){
 449                 ny = y - 7 - Dy(lr);
 450                 lr = rectsubpt(lr, Pt(0, lr.min.y - ny));
 451         }
 452         if(rectinrect(lr, screen->r) && (lr.min.x > x || lr.max.x <= x)){
 453                 string(screen, lr.min, display->black, ZP, display->defaultfont, buf);
 454                 return 1;
 455         }
 456         return 0;
 457 }
 458
 459 int
 460 yticklabel(char *fmt, double g, int p, int x, int y)
 461 {
 462         char buf[32];
 463         Rectangle lr;
 464         int nx;
 465
 466         snprint(buf, sizeof(buf), fmt, g);
 467         lr.min = Pt(0, 0);
 468         lr.max = stringsize(display->defaultfont, buf);
 469         lr = rectaddpt(lr, Pt(x-Dx(lr)-2, p - Dy(lr) / 2));
 470         if(lr.min.x < screen->r.min.x){
 471                 nx = x + 7;
 472                 lr = rectsubpt(lr, Pt(lr.min.x - nx, 0));
 473         }
 474         if(rectinrect(lr, screen->r) && (lr.min.y > y || lr.max.y <= y)){
 475                 string(screen, lr.min, display->black, ZP, display->defaultfont, buf);
 476                 return 1;
 477         }
 478         return 0;
 479 }
 480
 481 int
 482 calcm(double min, double max, int e, double *dp, double *mp)
 483 {
 484         double d, m, r;
 485
 486         d = pow(10, e>>1);
 487         if((e & 1) != 0) d *= 5;
 488         m = min;
 489         if(min < 0 && max > 0)
 490                 m += fmod(-m, d);
 491         else{
 492                 r = fmod(m, d);
 493                 if(r < 0)
 494                         m -= r;
 495                 else
 496                         m += d - r;
 497         }
 498         if(dp != nil) *dp = d;
 499         if(mp != nil) *mp = m;
 500         return (max-m)*0.999/d;
 501 }
 502
 503 int
 504 ticks(double min, double max, double *dp, double *mp)
 505 {
 506         int e, n;
 507         double m;
 508         int beste;
 509         double bestm;
 510
 511         e = 2 * ceil(log10(max - min));
 512         beste = 0;
 513         bestm = Inf(1);
 514         for(;e>-100;e--){
 515                 n = calcm(min, max, e, nil, nil);
 516                 if(n <= 0) continue;
 517                 if(n < 10) m = 10.0 / n;
 518                 else m = n / 10.0;
 519                 if(m < bestm){
 520                         beste = e;
 521                         bestm = m;
 522                 }
 523                 if(n > 10) break;
 524         }
 525         calcm(min, max, beste, dp, mp);
 526         return (max - *mp) / *dp;
 527 }
 528
 529 void
 530 drawaxes(void)
 531 {
 532         int x, y, p;
 533         double dx, dy, mx, my;
 534         int nx, ny;
 535         int i;
 536         char fmt[32];
 537
 538         if(xmin < 0 && xmax > 0)
 539                 x = deconvx(&screen->r, 0);
 540         else
 541                 x = screen->r.min.x+5;
 542         line(screen, Pt(x, screen->r.min.y), Pt(x, screen->r.max.y), Endarrow, 0, 0, display->black, ZP);
 543         if(ymin < 0 && ymax > 0)
 544                 y = deconvy(&screen->r, 0);
 545         else
 546                 y = screen->r.max.y-5;
 547         line(screen, Pt(screen->r.min.x, y), Pt(screen->r.max.x, y), 0, Endarrow, 0, display->black, ZP);
 548         nx = ticks(xmin, xmax, &dx, &mx);
 549         tickfmt(dx, mx, nx, fmt);
 550         for(i = 0; i <= nx; i++){
 551                 p = deconvx(&screen->r, dx*i+mx);
 552                 if(xticklabel(fmt, dx*i+mx, p, x, y))
 553                         line(screen, Pt(p, y), Pt(p, y-5), 0, 0, 0, display->black, ZP);
 554         }
 555         ny = ticks(ymin, ymax, &dy, &my);
 556         tickfmt(dy, my, ny, fmt);
 557         for(i = 0; i <= ny; i++){
 558                 p = deconvy(&screen->r, dy*i+my);
 559                 if(yticklabel(fmt, dy*i+my, p, x, y))
 560                         line(screen, Pt(x, p), Pt(x+5, p), 0, 0, 0, display->black, ZP);
 561         }
 562 }
 563
 564 void
 565 drawgraphs(void)
 566 {
 567         int i;
 568
 569         color = display->black;
 570         for(i = 0; i < nfns; i++)
 571                 drawgraph(&fns[i], &screen->r);
 572         if(!aflag)
 573                 drawaxes();
 574         flushimage(display, 1);
 575 }
 576
 577 void
 578 usage(void)
 579 {
 580         fprint(2, "usage: fplot [-a] [-c [-s size]] [-r range] functions ...\n");
 581         exits("usage");
 582 }
 583
 584 void
 585 zoom(void)
 586 {
 587         Mouse m;
 588         Rectangle r;
 589         double xmin_, xmax_, ymin_, ymax_;
 590
 591         m.buttons = 0;
 592         r = egetrect(1, &m);
 593         if(Dx(r) < 1 || Dy(r) < 1)
 594                 return;
 595         zoomst = realloc(zoomst, sizeof(FRectangle) * (nzoomst + 1));
 596         if(zoomst == nil) sysfatal("realloc: %r");
 597         zoomst[nzoomst++] = (FRectangle){xmin, xmax, ymin, ymax};
 598         xmin_ = convx(&screen->r, r.min.x);
 599         xmax_ = convx(&screen->r, r.max.x);
 600         ymin_ = convy(&screen->r, r.max.y);
 601         ymax_ = convy(&screen->r, r.min.y);
 602         xmin = xmin_;
 603         xmax = xmax_;
 604         ymin = ymin_;
 605         ymax = ymax_;
 606         draw(screen, screen->r, display->white, nil, ZP);
 607         drawgraphs();
 608 }
 609
 610 void
 611 unzoom(void)
 612 {
 613         if(nzoomst == 0) return;
 614         xmin = zoomst[nzoomst - 1].xmin;
 615         xmax = zoomst[nzoomst - 1].xmax;
 616         ymin = zoomst[nzoomst - 1].ymin;
 617         ymax = zoomst[nzoomst - 1].ymax;
 618         zoomst = realloc(zoomst, sizeof(FRectangle) * --nzoomst);
 619         if(zoomst == nil && nzoomst != 0) sysfatal("realloc: %r");
 620         draw(screen, screen->r, display->white, nil, ZP);
 621         drawgraphs();
 622 }
 623
 624 void
 625 parsefns(int n, char **s)
 626 {
 627         int i, max, cur;
 628
 629         max = 0;
 630         nfns = n;
 631         fns = emalloc(sizeof(*fns) * n);
 632         for(i = 0; i < nfns; i++) {
 633                 parse(&fns[i], s[i]);
 634                 cur = calcstack(&fns[i]);
 635                 if(cur > max)
 636                         max = cur;
 637         }
 638         stack = emalloc(sizeof(*stack) * max);
 639 }
 640
 641 void
 642 parserange(char *s)
 643 {
 644         while(*s && !isdigit(*s) && *s != '-') s++;
 645         if(*s == 0) return;
 646         xmin = strtod(s, &s);
 647         while(*s && !isdigit(*s) && *s != '-') s++;
 648         if(*s == 0) return;
 649         xmax = strtod(s, &s);
 650         while(*s && !isdigit(*s) && *s != '-') s++;
 651         if(*s == 0) return;
 652         ymin = strtod(s, &s);
 653         while(*s && !isdigit(*s) && *s != '-') s++;
 654         if(*s == 0) return;
 655         ymax = strtod(s, &s);
 656 }
 657
 658 void
 659 parsesize(char *s)
 660 {
 661         while(*s && !isdigit(*s)) s++;
 662         if(*s == 0) return;
 663         picx = strtol(s, &s, 0);
 664         while(*s && !isdigit(*s)) s++;
 665         if(*s == 0) return;
 666         picy = strtol(s, &s, 0);
 667 }
 668
 669 void
 670 main(int argc, char **argv)
 671 {
 672         Event e;
 673         Rectangle r;
 674         int i;
 675         static int lbut;
 676
 677         ARGBEGIN {
 678         case 'a': aflag++; break;
 679         case 'r': parserange(EARGF(usage())); break;
 680         case 's': parsesize(EARGF(usage())); break;
 681         case 'c': cflag++; break;
 682         default: usage();
 683         } ARGEND;
 684         if(argc < 1)
 685                 usage();
 686         setfcr(getfcr() & ~(FPZDIV | FPINVAL));
 687         parsefns(argc, argv);
 688         if(cflag) {
 689                 imagedata = emalloc(picx * picy * 3);
 690                 memset(imagedata, 0xFF, picx * picy * 3);
 691                 print("%11s %11d %11d %11d %11d ", "r8g8b8", 0, 0, picx, picy);
 692                 r.min.x = r.min.y = 0;
 693                 r.max.x = picx;
 694                 r.max.y = picy;
 695                 for(i = 0; i < nfns; i++)
 696                         drawgraph(&fns[i], &r);
 697                 if(write(1, imagedata, picx * picy * 3) < picx * picy * 3)
 698                         sysfatal("write: %r");
 699         } else {
 700                 if(initdraw(nil, nil, "fplot") < 0)
 701                         sysfatal("initdraw: %r");
 702                 einit(Emouse | Ekeyboard);
 703                 drawgraphs();
 704                 for(;;) {
 705                         switch(event(&e)) {
 706                         case Emouse:
 707                                 if((e.mouse.buttons & 1) != 0)
 708                                         zoom();
 709                                 if((~e.mouse.buttons & lbut & 4) != 0)
 710                                         unzoom();
 711                                 lbut = e.mouse.buttons;
 712                                 break;
 713                         case Ekeyboard:
 714                                 switch(e.kbdc) {
 715                                 case 'q': case 127: exits(nil); break;
 716                                 case 'y':
 717                                         if(!isInf(ymin, 1) && !isInf(ymax, -1)){
 718                                                 zoomst = realloc(zoomst, sizeof(FRectangle) * (nzoomst + 1));
 719                                                 if(zoomst == nil) sysfatal("realloc: %r");
 720                                                 zoomst[nzoomst++] = (FRectangle){xmin, xmax, ymin, ymax};
 721                                                 ymin = gymin-0.05*(gymax-gymin);
 722                                                 ymax = gymax+0.05*(gymax-gymin);
 723                                                 draw(screen, screen->r, display->white, nil, ZP);
 724                                                 drawgraphs();
 725                                         }
 726                                         break;
 727                                 }
 728                         }
 729                 }
 730         }
 731 }
 732
 733 void
 734 eresized(int new)
 735 {
 736         if(new) {
 737                 if(getwindow(display, Refnone) < 0)
 738                         sysfatal("getwindow: %r");
 739                 drawgraphs();
 740         }
 741 }