shrimp/src/prelude/index.ts

// The prelude creates all the builtin Shrimp functions.

import {
  type Value, type VM, toValue,
  extractParamInfo, isWrapped, getOriginalFunction,
} from 'reefvm'

import { dict } from './dict'
import { json } from './json'
import { load } from './load'
import { list } from './list'
import { math } from './math'
import { str } from './str'

export const globals = {
  dict,
  json,
  load,
  list,
  math,
  str,

  // hello
  echo: (...args: any[]) => {
    console.log(...args.map(a => {
      const v = toValue(a)
      return ['array', 'dict'].includes(v.type) ? formatValue(v, true) : v.value
    }))
    return toValue(null)
  },

  // info
  type: (v: any) => toValue(v).type,
  inspect: (v: any) => formatValue(toValue(v)),
  describe: (v: any) => {
    const val = toValue(v)
    return `#<${val.type}: ${formatValue(val)}>`
  },
  var: function (this: VM, v: any) {
    return typeof v === 'string' ? this.scope.get(v) : v
  },
  'var?': function (this: VM, v: string) {
    return typeof v !== 'string' || this.scope.has(v)
  },
  ref: (fn: Function) => fn,
  import: function (this: VM, atNamed: Record<any, string | string[]> = {}, ...idents: string[]) {
    const onlyArray = Array.isArray(atNamed.only) ? atNamed.only : [atNamed.only].filter(a => a)
    const only = new Set(onlyArray)
    const wantsOnly = only.size > 0


    for (const ident of idents) {
      const module = this.get(ident)

      if (!module) throw new Error(`import: can't find ${ident}`)
      if (module.type !== 'dict') throw new Error(`import: can't import ${module.type}`)

      for (const [name, value] of module.value.entries()) {
        if (value.type === 'dict') throw new Error(`import: can't import dicts in dicts`)
        if (wantsOnly && !only.has(name)) continue
        this.set(name, value)
      }
    }
  },

  // env
  args: Bun.argv.slice(1),
  exit: (num: number) => process.exit(num ?? 0),

  // type predicates
  'string?': (v: any) => toValue(v).type === 'string',
  'number?': (v: any) => toValue(v).type === 'number',
  'boolean?': (v: any) => toValue(v).type === 'boolean',
  'array?': (v: any) => toValue(v).type === 'array',
  'dict?': (v: any) => toValue(v).type === 'dict',
  'function?': (v: any) => {
    const t = toValue(v).type
    return t === 'function' || t === 'native'
  },
  'null?': (v: any) => toValue(v).type === 'null',
  'some?': (v: any) => toValue(v).type !== 'null',

  // boolean/logic
  not: (v: any) => !v,
  bnot: (n: number) => ~(n | 0),

  // utilities
  inc: (n: number) => n + 1,
  dec: (n: number) => n - 1,
  identity: (v: any) => v,

  // collections
  length: (v: any) => {
    const value = toValue(v)
    switch (value.type) {
      case 'string': case 'array': return value.value.length
      case 'dict': return value.value.size
      default: throw new Error(`length: expected string, array, or dict, got ${value.type}`)
    }
  },
  at: (collection: any, index: number | string) => {
    const value = toValue(collection)
    if (value.type === 'string' || value.type === 'array') {
      const idx = typeof index === 'number' ? index : parseInt(index as string)
      if (idx < 0 || idx >= value.value.length) {
        throw new Error(`at: index ${idx} out of bounds for ${value.type} of length ${value.value.length}`)
      }
      return value.value[idx]
    } else if (value.type === 'dict') {
      const key = String(index)
      if (!value.value.has(key)) {
        throw new Error(`at: key '${key}' not found in dict`)
      }
      return value.value.get(key)
    } else {
      throw new Error(`at: expected string, array, or dict, got ${value.type}`)
    }
  },
  range: (start: number, end: number | null) => {
    if (end === null) {
      end = start
      start = 0
    }
    const result: number[] = []
    for (let i = start; i <= end; i++) {
      result.push(i)
    }
    return result
  },
  'empty?': (v: any) => {
    const value = toValue(v)
    switch (value.type) {
      case 'string': case 'array':
        return value.value.length === 0
      case 'dict':
        return value.value.size === 0
      default:
        return false
    }
  },

  // enumerables
  each: async (list: any[], cb: Function) => {
    for (const value of list) await cb(value)
    return list
  },

}

export const colors = {
  reset: '\x1b[0m',
  bright: '\x1b[1m',
  dim: '\x1b[2m',
  cyan: '\x1b[36m',
  yellow: '\x1b[33m',
  green: '\x1b[32m',
  red: '\x1b[31m',
  blue: '\x1b[34m',
  magenta: '\x1b[35m',
  pink: '\x1b[38;2;255;105;180m'
}

export function formatValue(value: Value, inner = false): string {
  switch (value.type) {
    case 'string':
      return `${colors.green}'${value.value.replaceAll("'", "\\'")}${colors.green}'${colors.reset}`
    case 'number':
      return `${colors.cyan}${value.value}${colors.reset}`
    case 'boolean':
      return `${colors.yellow}${value.value}${colors.reset}`
    case 'null':
      return `${colors.dim}null${colors.reset}`
    case 'array': {
      const items = value.value.map(x => formatValue(x, true)).join(' ')
      return `${colors.blue}[${colors.reset}${items}${colors.blue}]${colors.reset}`
    }
    case 'dict': {
      const entries = Array.from(value.value.entries())
        .map(([k, v]) => `${k}${colors.blue}=${colors.reset}${formatValue(v, true)}`)
        .join(' ')
      if (entries.length === 0)
        return `${colors.blue}[=]${colors.reset}`
      return `${colors.blue}[${colors.reset}${entries}${colors.blue}]${colors.reset}`
    }
    case 'function': {
      const params = value.params.length ? '(' + value.params.join(' ') + ')' : ''
      return `${colors.dim}<function${params}>${colors.reset}`
    }
    case 'native':
      const fn = isWrapped(value.fn) ? getOriginalFunction(value.fn) : value.fn
      const info = extractParamInfo(fn)
      const params = info.params.length ? '(' + info.params.join(' ') + ')' : ''
      return `${colors.dim}<native${params}>${colors.reset}`
    case 'regex':
      return `${colors.magenta}${value.value}${colors.reset}`
    default:
      return String(value)
  }
}