mlish: add nbody test

tapl/mlish.rkt

@@ -14,7 +14,7 @@
 (provide → × tup proj define-type-alias)
 (provide define-type match)
 (provide (rename-out [ext-stlc:and and] [ext-stlc:#%datum #%datum]))
-(reuse [cons stlc:cons] nil isnil head tail [list stlc:list] List ~List List? #:from "stlc+cons.rkt")
+(reuse reverse [cons stlc:cons] nil isnil head tail [list stlc:list] List ~List List? #:from "stlc+cons.rkt")
 (provide (rename-out [stlc:list list] [stlc:cons cons]))
 (reuse ref deref := Ref #:from "stlc+box.rkt")
 
@@ -752,12 +752,15 @@
    (⊢ (string-copy! dest- dest-start- src- src-start- src-end-) : Unit)])
 
 (define-primop fl+ : (→ Float Float Float))
+(define-primop fl- : (→ Float Float Float))
 (define-primop fl* : (→ Float Float Float))
 (define-primop fl/ : (→ Float Float Float))
+(define-primop flsqrt : (→ Float Float))
 (define-primop flceiling : (→ Float Float))
 (define-primop inexact->exact : (→ Float Int))
 (define-primop exact->inexact : (→ Int Float))
 (define-primop char->integer : (→ Char Int))
+(define-primop real->decimal-string : (→ Float Int String))
 (define-primop fx->fl : (→ Int Float))
 (define-typed-syntax quotient+remainder
   [(_ x y)

tapl/stlc+cons.rkt

@@ -60,4 +60,9 @@
    #:with (~List τ) (local-expand #'expected-τ 'expression null)
    #'(cons/tc (add-expected x τ) (list/tc . rst))]
   [(_ x . rst) ; no expected type
-   #'(cons/tc x (list/tc . rst))])
+   #'(cons/tc x (list/tc . rst))])
+(define-typed-syntax reverse
+  [(_ e)
+   #:with (e- τ-lst) (infer+erase #'e)
+   #:when (List? #'τ-lst)
+   (⊢ (reverse e-) : τ-lst)])

tapl/tests/mlish/nbody.mlish

@@ -0,0 +1,185 @@
+#lang s-exp "../../mlish.rkt"
+(require "../rackunit-typechecking.rkt")
+
+(define +pi+ 3.141592653589793)
+
+(check-type +pi+ : Float)
+
+(define +days-per-year+ 365.24)
+
+(define * fl*)
+
+(define +solar-mass+ (* 4.0 (* +pi+ +pi+)))
+
+(check-type +solar-mass+ : Float)
+
+(define +dt+ 0.01)
+
+(define-type Body (body Float ; x
+                        Float ; y
+                        Float ; z
+                        Float ; vx
+                        Float ; vy
+                        Float ; vz
+                        Float ; mass
+                        ))
+
+(define *sun*
+  (body 0.0 0.0 0.0 0.0 0.0 0.0 +solar-mass+))
+
+(define *jupiter*
+  (body 4.84143144246472090
+             -1.16032004402742839
+             -1.03622044471123109e-1
+             (* 1.66007664274403694e-3 +days-per-year+)
+             (* 7.69901118419740425e-3 +days-per-year+)
+             (* -6.90460016972063023e-5 +days-per-year+)
+             (* 9.54791938424326609e-4 +solar-mass+)))
+
+(define *saturn*
+  (body 8.34336671824457987
+             4.12479856412430479
+             -4.03523417114321381e-1
+             (* -2.76742510726862411e-3 +days-per-year+)
+             (* 4.99852801234917238e-3 +days-per-year+)
+             (* 2.30417297573763929e-5 +days-per-year+)
+             (* 2.85885980666130812e-4 +solar-mass+)))
+
+(define *uranus*
+  (body 1.28943695621391310e1
+             -1.51111514016986312e1
+             -2.23307578892655734e-1
+             (* 2.96460137564761618e-03 +days-per-year+)
+             (* 2.37847173959480950e-03 +days-per-year+)
+             (* -2.96589568540237556e-05 +days-per-year+)
+             (*  4.36624404335156298e-05 +solar-mass+)))
+
+(define *neptune*
+  (body 1.53796971148509165e+01
+             -2.59193146099879641e+01
+             1.79258772950371181e-01
+             (* 2.68067772490389322e-03 +days-per-year+)
+             (* 1.62824170038242295e-03 +days-per-year+)
+             (* -9.51592254519715870e-05 +days-per-year+)
+             (* 5.15138902046611451e-05 +solar-mass+)))
+
+(define *system* (list *sun* *jupiter* *saturn* *uranus* *neptune*))
+
+(check-type *system* : (List Body))
+
+(define (offset-momentum -> Unit)
+  (let loop-i : Unit 
+    ([i *system*] [px 0.0] [py 0.0] [pz 0.0])
+    (cond
+     [(isnil i)
+      (match (head *system*) with ; sun
+       [body x y z vx vy vz m ->
+        (let ([new 
+               (body x y z
+                     (fl/ (fl- 0.0 px) +solar-mass+)
+                     (fl/ (fl- 0.0 py) +solar-mass+)
+                     (fl/ (fl- 0.0 pz) +solar-mass+)
+                     m)])
+          (set! *system* (cons new (tail *system*))))])]
+     [else
+      (match (head i) with
+       [body x y z vx vy vz m ->
+        (loop-i (tail i)
+          (fl+ px (fl* vx m))
+          (fl+ py (fl* vy m))
+          (fl+ pz (fl* vz m)))])])))
+
+(define (energy [o : (List Body)] -> Float)
+  (let loop-o : Float ([o o] [e 0.0])
+    (cond 
+     [(isnil o) e]
+     [else
+      (match (head o) with
+       [body x y z vx vy vz m ->
+        (let* ([e (fl+ e (fl* 0.5 
+                         (fl* m
+                           (fl+ (fl+ (fl* vx vx)
+                                  (fl* vy vy))
+                             (fl* vz vz)))))])
+        (let loop-i : Float ([i (tail o)] [e e])
+          (if (isnil i)
+              (loop-o (tail o) e)
+              (match (head i) with
+               [body x2 y2 z2 vx2 vy2 vz2 m2 ->
+                (let* ([dx   (fl- x x2)]
+                       [dy   (fl- y y2)]
+                       [dz   (fl- z z2)]
+                       [dist (flsqrt (fl+ (fl+ (fl* dx dx) (fl* dy dy)) 
+                                       (fl* dz dz)))]
+                       [e    (fl- e (fl/ (fl* m m2) dist))])
+              (loop-i (tail i) e))]))))])])))
+
+(define (advance [bs : (List Body)] -> (List Body))
+  (if (isnil bs)
+      bs
+      (let ([new (advance2 bs)])
+        (cons (head new) (advance (tail new))))))
+;; bs is non-null
+(define (advance2 [bs : (List Body)] -> (List Body))
+  (match (head bs) with
+   [body o1x o1y o1z vx vy vz om ->
+    (let loop-i : (List Body) 
+      ([i (tail bs)] [res (nil {Body})] [vx vx] [vy vy] [vz vz])
+      (cond 
+       [(isnil i) 
+        (cons 
+          (body 
+           (fl+ o1x (fl* +dt+ vx))
+           (fl+ o1y (fl* +dt+ vy))
+           (fl+ o1z (fl* +dt+ vz))
+           vx vy vz om)
+          (reverse res))]
+       [else
+        (match (head i) with
+         [body i1x i1y i1z i1vx i1vy i1vz im ->
+          (let* ([dx    (fl- o1x i1x)]
+                 [dy    (fl- o1y i1y)]
+                 [dz    (fl- o1z i1z)]
+                 [dist2 (fl+ (fl+ (fl* dx dx) (fl* dy dy)) (fl* dz dz))]
+                 [mag   (fl/ +dt+ (fl* dist2 (flsqrt dist2)))]
+                 [dxmag (fl* dx mag)]
+                 [dymag (fl* dy mag)]
+                 [dzmag (fl* dz mag)])
+            (loop-i 
+              (tail i)
+              (cons (body i1x i1y i1z
+                          (fl+ i1vx (fl* dxmag om))
+                          (fl+ i1vy (fl* dymag om))
+                          (fl+ i1vz (fl* dzmag om))
+                          im) res)
+              (fl- vx (fl* dxmag im))
+              (fl- vy (fl* dymag im))
+              (fl- vz (fl* dzmag im))))])]))]))
+    
+(check-type (real->decimal-string (energy *system*) 9) 
+  : String -> "-0.169289903")
+
+(offset-momentum)
+
+(check-type (real->decimal-string (energy *system*) 9) 
+  : String -> "-0.169075164")
+
+(check-type 
+    (real->decimal-string 
+        (energy (advance *system*)) 
+      9)
+  : String -> "-0.169074954")
+
+(check-type 
+    (real->decimal-string 
+        (energy (advance (advance *system*))) 9)
+  : String -> "-0.169074743")
+
+(check-type
+    (real->decimal-string
+        (energy
+            (for/fold ([bs *system*]) 
+                      ([i (in-range 10)]) 
+              (advance bs)))
+      9)
+  : String -> "-0.169073022")

tapl/tests/run-all-mlish-tests.rkt

@@ -12,3 +12,4 @@
 ;(require "mlish/heapsort.mlish")
 (require "mlish/knuc.mlish")
 (require "mlish/matrix.mlish")
+(require "mlish/nbody.mlish")