nanomath/src/complex/Complex.js

import {Returns, Type} from '#core/Type.js'
import {match} from '#core/TypePatterns.js'

const isComplex = z => z && typeof z === 'object' && 're' in z && 'im' in z

const specializesTo = CType => CType.complex

function complexSpecialize(ComponentType) {
   const compTest = ComponentType.test
   const specTest = z => isComplex(z) && compTest(z.re) && compTest(z.im)
   const typeName = `Complex(${ComponentType})`
   if (ComponentType.concrete) {
      const fromSpec = [match(
         ComponentType, () => r => ({re: r, im: ComponentType.zero}))]
      for (const {pattern, does} of ComponentType.from) {
         fromSpec.push(match(
            this.specialize(pattern.type),
            math => {
               const compConv = does(math)
               return z => ({re: compConv(z.re), im: compConv(z.im)})
            }))
      }
      const typeOptions = {from: fromSpec, typeName}
      if ('zero' in ComponentType) {
         typeOptions.zero = {re: ComponentType.zero, im: ComponentType.zero}
         if ('one' in ComponentType) {
            typeOptions.one = {re: ComponentType.one, im: ComponentType.zero}
         }
      }
      if ('nan' in ComponentType) {
         typeOptions.nan = {re: ComponentType.nan, im: ComponentType.nan}
      }
      const complexCompType = new Type(specTest, typeOptions)
      complexCompType.Component = ComponentType
      complexCompType.complex = true
      return complexCompType
   }
   // wrapping a generic type in Complex
   const cplx = this

   const innerSpecialize = (...args) => {
      const innerType = ComponentType.specialize(...args)
      return cplx.specialize(innerType)
   }

   const innerSpecializesTo = CType => specializesTo(CType)
       && ComponentType.specializesTo(CType.Component)

   const innerRefine = (z, typer) => {
      const reType = ComponentType.refine(z.re, typer)
      const imType = ComponentType.refine(z.im, typer)
      if (reType === imType) return cplx.specialize(reType)
      throw new TypeError(
         'mixed-type Complex numbers disallowed '
            + `(real: ${reType}, imaginary: ${imType})`)
   }

   return new Type(specTest, {
      specialize: innerSpecialize,
      specializesTo: innerSpecializesTo,
      refine: innerRefine,
      typeName,
   })
}

export const Complex = new Type(isComplex, {
   specialize: complexSpecialize,
   specializesTo,
   from: [match( // can promote any non-complex type T to Complex(T) as needed
      // but watch out, this should be tried late, because it can preclude
      // other more reasonable conversions like bool => number.
      T => !T.complex,
      (math, T) => Returns(Complex(T), r => math.complex(r, math.zero(T)))
   )],
   refine: function(z, typer) {
      const reType = typer(z.re)
      const imType = typer(z.im)
      if (reType === imType) return this.specialize(reType)
      throw new TypeError(
         'mixed-type Complex numbers disallowed '
         + `(real: ${reType}, imaginary: ${imType})`)
   }
})