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