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)
  },

  // 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,

  // 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)
  }
}