feat: Add return typing strategies and implement sqrt with them (#26)

Resolves #25

Reviewed-on: #26
Co-authored-by: Glen Whitney <glen@studioinfinity.org>
Co-committed-by: Glen Whitney <glen@studioinfinity.org>

src/complex/arithmetic.js

@@ -1,43 +1,46 @@
 import {Complex} from './Complex.js'
-import {promoteBinary, promoteUnary} from './helpers.js'
+import {maybeComplex, promoteBinary, promoteUnary} from './helpers.js'
 import {ResolutionError} from '#core/helpers.js'
+import {Returns, ReturnTyping} from '#core/Type.js'
 import {match} from '#core/TypePatterns.js'
 import {ReturnsAs} from '#generic/helpers.js'
 
-export const absquare = match(Complex, (math, C) => {
-   const compAbsq = math.absquare.resolve([C.Component])
+const {conservative, full, free} = ReturnTyping
+
+export const absquare = match(Complex, (math, C, strategy) => {
+   const compAbsq = math.absquare.resolve(C.Component, full)
    const R = compAbsq.returns
-   const add = math.add.resolve([R,R])
+   const add = math.add.resolve([R,R], strategy)
    return ReturnsAs(add, z => add(compAbsq(z.re), compAbsq(z.im)))
 })
 
 export const add = promoteBinary('add')
 
-export const conj = match(Complex, (math, C) => {
-   const neg = math.negate.resolve(C.Component)
-   const compConj = math.conj.resolve(C.Component)
-   const cplx = math.complex.resolve([compConj.returns, neg.returns])
+export const conj = match(Complex, (math, C, strategy) => {
+   const neg = math.negate.resolve(C.Component, full)
+   const compConj = math.conj.resolve(C.Component, full)
+   const cplx = maybeComplex(math, strategy, compConj.returns, neg.returns)
    return ReturnsAs(cplx, z => cplx(compConj(z.re), neg(z.im)))
 })
 
 export const divide = [
-   match([Complex, T => !T.complex], (math, [C, R]) => {
-      const div = math.divide.resolve([C.Component, R])
-      const cplx = math.complex.resolve([div.returns, div.returns])
+   match([Complex, T => !T.complex], (math, [C, R], strategy) => {
+      const div = math.divide.resolve([C.Component, R], full)
+      const cplx = maybeComplex(math, strategy, div.returns, div.returns)
       return ReturnsAs(cplx, (z, r) => cplx(div(z.re, r), div(z.im, r)))
    }),
-   match([Complex, Complex], (math, [W, Z]) => {
-      const inv = math.invert.resolve(Z)
-      const mult = math.multiply.resolve([W, inv.returns])
+   match([Complex, Complex], (math, [W, Z], strategy) => {
+      const inv = math.invert.resolve(Z, full)
+      const mult = math.multiply.resolve([W, inv.returns], strategy)
       return ReturnsAs(mult, (w, z) => mult(w, inv(z)))
    })
 ]
 
-export const invert = match(Complex, (math, C) => {
-   const conj = math.conj.resolve(C)
-   const norm = math.absquare.resolve(C)
-   const div = math.divide.resolve([C.Component, norm.returns])
-   const cplx = math.complex.resolve([div.returns, div.returns])
+export const invert = match(Complex, (math, C, strategy) => {
+   const conj = math.conj.resolve(C, full)
+   const norm = math.absquare.resolve(C, full)
+   const div = math.divide.resolve([C.Component, norm.returns], full)
+   const cplx = maybeComplex(math, strategy, div.returns, div.returns)
    return ReturnsAs(cplx, z => {
       const c = conj(z)
       const d = norm(z)
@@ -47,23 +50,87 @@ export const invert = match(Complex, (math, C) => {
 
 // We want this to work for complex numbers, quaternions, octonions, etc
 // See https://math.ucr.edu/home/baez/octonions/node5.html
-export const multiply = match([Complex, Complex], (math, [W, Z]) => {
-   const conj = math.conj.resolve(W.Component)
-   if (conj.returns !== W.Component) {
-      throw new ResolutionError(
-         `conjugation on ${W.Component} returns other type (${conj.returns})`)
-   }
-   const mWZ = math.multiply.resolve([W.Component, Z.Component])
-   const mZW = math.multiply.resolve([Z.Component, W.Component])
-   const sub = math.subtract.resolve([mWZ.returns, mZW.returns])
-   const add = math.add.resolve([mWZ.returns, mZW.returns])
-   const cplx = math.complex.resolve([sub.returns, add.returns])
-   return ReturnsAs(cplx, (w, z) => {
-      const real = sub(mWZ(     w.re,  z.re), mZW(z.im, conj(w.im)))
-      const imag = add(mWZ(conj(w.re), z.im), mZW(z.re,      w.im))
-      return cplx(real, imag)
+export const multiply = [
+   match([T => !T.complex, Complex], (math, [R, C], strategy) => {
+      const mult = math.multiply.resolve([R, C.Component], full)
+      const cplx = maybeComplex(math, strategy, mult.returns, mult.returns)
+      return ReturnsAs(cplx, (r, z) => cplx(mult(r, z.re), mult(r, z.im)))
+   }),
+   match([Complex, T => !T.complex], (math, [C, R], strategy) => {
+      const mult = math.multiply.resolve([R, C], strategy)
+      return ReturnsAs(mult, (z, r) => mult(r, z))
+   }),
+   match([Complex, Complex], (math, [W, Z], strategy) => {
+      const conj = math.conj.resolve(W.Component, full)
+      if (conj.returns !== W.Component) {
+         throw new ResolutionError(
+            `conjugation on ${W.Component} returns type (${conj.returns})`)
+      }
+      const mWZ = math.multiply.resolve([W.Component, Z.Component], full)
+      const mZW = math.multiply.resolve([Z.Component, W.Component], full)
+      const sub = math.subtract.resolve([mWZ.returns, mZW.returns], full)
+      const add = math.add.resolve([mWZ.returns, mZW.returns], full)
+      const cplx = maybeComplex(math, strategy, sub.returns, add.returns)
+      return ReturnsAs(cplx, (w, z) => {
+         const real = sub(mWZ(     w.re,  z.re), mZW(z.im, conj(w.im)))
+         const imag = add(mWZ(conj(w.re), z.im), mZW(z.re,      w.im))
+         return cplx(real, imag)
+      })
    })
-})
+]
 
 export const negate = promoteUnary('negate')
+
+// Should work for complex, quaternions, octonions, etc., even with
+// integer coordinates.
+export const sqrt = match(Complex, (math, C, strategy) => {
+   const re = math.re.resolve(C)
+   const R = re.returns
+   const isReal = math.isReal.resolve(C)
+   // dependencies for the real case:
+   const zComp = math.zero(C.Component)
+   const sign = math.sign.resolve(R, conservative)
+   const oneR = math.one(R)
+   const zR = math.zero(R)
+   const addRComp = math.add.resolve([R, C.Component], full)
+   const sqrtR = math.sqrt.resolve(R, conservative)
+   const neg = math.negate.resolve(R, conservative)
+   const cplx = math.complex.resolve([C.Component, C.Component], full)
+   // additional dependencies for the complex case
+   const abs = math.abs.resolve(C, full)
+   if (abs.returns !== R) {
+      throw new TypeError(`abs on ${C} returns ${abs.returns}, not ${R}`)
+   }
+   const addRR = math.add.resolve([R, R], conservative)
+   const twoR = addRR(oneR, oneR)
+   const multRR = math.multiply.resolve([R, R], conservative)
+   const im = math.im.resolve(C, full)
+   const addRC = math.add.resolve([R, C], full)
+   const divRR = math.divide.resolve([R, R], conservative)
+   const divCR = math.divide.resolve([C, R], full)
+
+   // The guts of the computation:
+   const sqrtImp = Returns(C, z => {
+      const a = re(z)
+      if (isReal(z)) { // always a special case
+         let rp = zComp
+         let ip = zComp
+         const sgn = sign(a)
+         if (sgn === oneR) rp = addRComp(sqrtR(a), rp)
+         else if (sgn !== zR) ip = addRComp(sqrtR(neg(a)), ip)
+         return cplx(rp, ip)
+      }
+      // Complex case:
+      // We can write z = a + q where q is pure imaginary.
+      // Let s = sqrt(2(|z|+a)). Then sqrt(z) = (s/2) + (q/s).
+      const s = sqrtR(multRR(twoR, addRR(abs(z), a)))
+      const q = im(z)
+      return addRC(divRR(s, twoR), divCR(q, s))
+   })
+
+   if (strategy != free) return sqrtImp
+   const prune = math.pruneImaginary.resolve(C, free)
+   return ReturnsAs(prune, z => prune(sqrtImp(z)))
+})
+
 export const subtract = promoteBinary('subtract')