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